From a4633f56774d91b7ead807b7c79883279e9a34b1 Mon Sep 17 00:00:00 2001
From: Denis Arh <denis.arh@gmail.com>
Date: Wed, 14 Apr 2021 17:24:25 +0200
Subject: [PATCH] Improve wf debugging capabilities (session tracing)

---
 automation/service/session.go | 33 ++++++++++++++++-----------------
 automation/types/session.go   | 11 ++++++++++-
 2 files changed, 26 insertions(+), 18 deletions(-)

diff --git a/automation/service/session.go b/automation/service/session.go
index 2d731bfab..2165084bf 100644
--- a/automation/service/session.go
+++ b/automation/service/session.go
@@ -179,14 +179,9 @@ func (svc *session) Start(g *wfexec.Graph, i auth.Identifiable, ssp types.Sessio
 
 	ses.CreatedAt = *now()
 	ses.CreatedBy = i.Identity()
+	ses.Status = types.SessionStarted
 	ses.Apply(ssp)
 
-	if ssp.Trace {
-		if err = store.CreateAutomationSession(context.TODO(), svc.store, ses); err != nil {
-			return
-		}
-	}
-
 	if err = ses.Exec(ctx, start, ssp.Input); err != nil {
 		return
 	}
@@ -344,21 +339,20 @@ func (svc *session) stateChangeHandler(ctx context.Context) wfexec.StateChangeHa
 		var (
 			// By default we want to update session when new status is prompted, delayed, completed or failed
 			//
-			// But if status is active, we
+			// But if status is active, we'll flush it every X frames (flushFrquency)
 			update = true
-			frame  = state.MakeFrame()
+
+			frame = state.MakeFrame()
 		)
 
 		// Stacktrace will be set to !nil if frame collection is needed
-		if ses.Stacktrace != nil {
-			if len(ses.Stacktrace) > 0 {
-				// calculate how long it took to get to this step
-				frame.ElapsedTime = uint(frame.CreatedAt.Sub(ses.Stacktrace[0].CreatedAt) / time.Millisecond)
-			}
-
-			ses.Stacktrace = append(ses.Stacktrace, frame)
+		if len(ses.RuntimeStacktrace) > 0 {
+			// calculate how long it took to get to this step
+			frame.ElapsedTime = uint(frame.CreatedAt.Sub(ses.RuntimeStacktrace[0].CreatedAt) / time.Millisecond)
 		}
 
+		ses.RuntimeStacktrace = append(ses.RuntimeStacktrace, frame)
+
 		switch i {
 		case wfexec.SessionPrompted:
 			ses.SuspendedAt = now()
@@ -381,14 +375,19 @@ func (svc *session) stateChangeHandler(ctx context.Context) wfexec.StateChangeHa
 
 		default:
 			// force update on every 10 new frames but only when stacktrace is not nil
-			update = ses.Stacktrace != nil && len(ses.Stacktrace)%flushFrequency == 0
+			update = ses.RuntimeStacktrace != nil && len(ses.RuntimeStacktrace)%flushFrequency == 0
 		}
 
 		if !update {
 			return
 		}
 
-		if err := svc.store.UpdateAutomationSession(ctx, ses); err != nil {
+		if ses.Stacktrace != nil || ses.Error != "" {
+			// Save stacktrace when we know we're tracing workflows OR whenever there is an error...
+			ses.Stacktrace = ses.RuntimeStacktrace
+		}
+
+		if err := svc.store.UpsertAutomationSession(ctx, ses); err != nil {
 			log.Error("failed to update session", zap.Error(err))
 		} else {
 			log.Debug("session updated", zap.Stringer("status", ses.Status))
diff --git a/automation/types/session.go b/automation/types/session.go
index dde5e006f..f07e845bc 100644
--- a/automation/types/session.go
+++ b/automation/types/session.go
@@ -25,6 +25,7 @@ type (
 		Input  *expr.Vars `json:"input"`
 		Output *expr.Vars `json:"output"`
 
+		// Stacktrace that gets stored (if/when configured)
 		Stacktrace Stacktrace `json:"stacktrace"`
 
 		CreatedAt time.Time  `json:"createdAt,omitempty"`
@@ -38,6 +39,13 @@ type (
 		Error       string     `json:"error,omitempty"`
 
 		session *wfexec.Session
+
+		// For keeping runtime stacktrace,
+		// even if we do not want to store it on every update
+		//
+		// This will aid us when session fails and we can access
+		// the whole stacktrace
+		RuntimeStacktrace Stacktrace `json:"-"`
 	}
 
 	SessionStartParams struct {
@@ -124,7 +132,8 @@ func (s *Session) Apply(ssp SessionStartParams) {
 	}
 
 	if ssp.Trace {
-		// set prop
+		// set Stacktrace prop to signal status handler
+		// that we're interested in storing stacktrace
 		s.Stacktrace = Stacktrace{}
 	}
 }