From be9e412e1cd9657ac183018059121d54ae6510d5 Mon Sep 17 00:00:00 2001
From: Mihail Stoykov <MStoykov@users.noreply.github.com>
Date: Mon, 2 Nov 2020 15:23:51 +0200
Subject: [PATCH] catching of VU panics (#1697)

fixes #1601
---
 js/runner.go      | 24 ++++++++++++++---
 js/runner_test.go | 65 +++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 85 insertions(+), 4 deletions(-)

diff --git a/js/runner.go b/js/runner.go
index 136c29d45db8..680d886128ce 100644
--- a/js/runner.go
+++ b/js/runner.go
@@ -21,6 +21,7 @@
 package js
 
 import (
+	"bytes"
 	"context"
 	"crypto/tls"
 	"encoding/json"
@@ -28,6 +29,7 @@ import (
 	"net"
 	"net/http"
 	"net/http/cookiejar"
+	"runtime/debug"
 	"strconv"
 	"time"
 
@@ -601,9 +603,8 @@ func (u *ActiveVU) RunOnce() error {
 
 func (u *VU) runFn(
 	ctx context.Context, isDefault bool, fn goja.Callable, args ...goja.Value,
-) (goja.Value, bool, time.Duration, error) {
+) (v goja.Value, isFullIteration bool, t time.Duration, err error) {
 	if !u.Runner.Bundle.Options.NoCookiesReset.ValueOrZero() {
-		var err error
 		u.state.CookieJar, err = cookiejar.New(nil)
 		if err != nil {
 			return goja.Undefined(), false, time.Duration(0), err
@@ -621,11 +622,26 @@ func (u *VU) runFn(
 	u.Runtime.Set("__ITER", u.Iteration)
 	u.Iteration++
 
+	defer func() {
+		if r := recover(); r != nil {
+			gojaStack := u.Runtime.CaptureCallStack(20, nil)
+			err = fmt.Errorf("a panic occurred in VU code but was caught: %s", r)
+			// TODO figure out how to use PanicLevel without panicing .. this might require changing
+			// the logger we use see
+			// https://github.com/sirupsen/logrus/issues/1028
+			// https://github.com/sirupsen/logrus/issues/993
+			b := new(bytes.Buffer)
+			for _, s := range gojaStack {
+				s.Write(b)
+			}
+			u.state.Logger.Log(logrus.ErrorLevel, "panic: ", r, "\n", string(debug.Stack()), "\nGoja stack:\n", b.String())
+		}
+	}()
+
 	startTime := time.Now()
-	v, err := fn(goja.Undefined(), args...) // Actually run the JS script
+	v, err = fn(goja.Undefined(), args...) // Actually run the JS script
 	endTime := time.Now()
 
-	var isFullIteration bool
 	select {
 	case <-ctx.Done():
 		isFullIteration = false
diff --git a/js/runner_test.go b/js/runner_test.go
index 30287f0dc180..fc9e9aa4984e 100644
--- a/js/runner_test.go
+++ b/js/runner_test.go
@@ -33,10 +33,12 @@ import (
 	"net/http"
 	"net/url"
 	"os"
+	"strings"
 	"sync"
 	"testing"
 	"time"
 
+	"github.com/sirupsen/logrus"
 	logtest "github.com/sirupsen/logrus/hooks/test"
 	"github.com/spf13/afero"
 	"github.com/stretchr/testify/assert"
@@ -1735,3 +1737,66 @@ func TestSystemTags(t *testing.T) {
 		})
 	}
 }
+
+func TestVUPanic(t *testing.T) {
+	r1, err := getSimpleRunner(t, "/script.js", `
+			var group = require("k6").group;
+			exports.default = function() {
+				group("panic here", function() {
+					if (__ITER == 0) {
+						panic("here we panic");
+					}
+					console.log("here we don't");
+				})
+			}`,
+	)
+	require.NoError(t, err)
+
+	r2, err := NewFromArchive(testutils.NewLogger(t), r1.MakeArchive(), lib.RuntimeOptions{})
+	if !assert.NoError(t, err) {
+		return
+	}
+
+	runners := map[string]*Runner{"Source": r1, "Archive": r2}
+	for name, r := range runners {
+		r := r
+		t.Run(name, func(t *testing.T) {
+			initVU, err := r.NewVU(1234, make(chan stats.SampleContainer, 100))
+			if !assert.NoError(t, err) {
+				return
+			}
+
+			logger := logrus.New()
+			logger.SetLevel(logrus.InfoLevel)
+			logger.Out = ioutil.Discard
+			hook := testutils.SimpleLogrusHook{
+				HookedLevels: []logrus.Level{logrus.InfoLevel, logrus.ErrorLevel, logrus.FatalLevel, logrus.PanicLevel},
+			}
+			logger.AddHook(&hook)
+
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+			vu := initVU.Activate(&lib.VUActivationParams{RunContext: ctx})
+			vu.(*ActiveVU).Runtime.Set("panic", func(str string) { panic(str) })
+			vu.(*ActiveVU).state.Logger = logger
+
+			vu.(*ActiveVU).Console.logger = logger.WithField("source", "console")
+			err = vu.RunOnce()
+			require.Error(t, err)
+			assert.Contains(t, err.Error(), "a panic occurred in VU code but was caught: here we panic")
+			entries := hook.Drain()
+			require.Len(t, entries, 1)
+			assert.Equal(t, logrus.ErrorLevel, entries[0].Level)
+			require.True(t, strings.HasPrefix(entries[0].Message, "panic: here we panic"))
+			require.True(t, strings.HasSuffix(entries[0].Message, "Goja stack:\nfile:///script.js:3:4(12)"))
+
+			err = vu.RunOnce()
+			assert.NoError(t, err)
+
+			entries = hook.Drain()
+			require.Len(t, entries, 1)
+			assert.Equal(t, logrus.InfoLevel, entries[0].Level)
+			require.Contains(t, entries[0].Message, "here we don't")
+		})
+	}
+}