From f17c2eb08a8c82ad1522221463c16a7312a0c715 Mon Sep 17 00:00:00 2001
From: samrabelachew <samrabelachew@users.noreply.github.com>
Date: Wed, 26 Jul 2023 12:12:51 -0700
Subject: [PATCH] Check if key is in roots map of Processor (#718)

Before notifying, we should ensure the key is in the roots map to
prevent attempts to notify check run on invalid, empty repo metadata.
There is a chance the root map can miss a key if multiple revisions come
in quickly and a new Processor for the latest revision is receiving a
signal on an old revision. We shouldn't care about this signal since we
already perform a clean defer to cancel those check runs on shutdown of
the old revision Processor.
---
 .../workflows/internal/pr/revision/state.go   | 18 ++++++++++++-
 .../internal/pr/revision/state_test.go        | 26 ++++++++++++++++++-
 2 files changed, 42 insertions(+), 2 deletions(-)

diff --git a/server/neptune/workflows/internal/pr/revision/state.go b/server/neptune/workflows/internal/pr/revision/state.go
index 71d9a34b0..633d7f15a 100644
--- a/server/neptune/workflows/internal/pr/revision/state.go
+++ b/server/neptune/workflows/internal/pr/revision/state.go
@@ -1,6 +1,7 @@
 package revision
 
 import (
+	"fmt"
 	"github.com/pkg/errors"
 	"github.com/runatlantis/atlantis/server/metrics"
 	"github.com/runatlantis/atlantis/server/neptune/workflows/internal/notifier"
@@ -9,6 +10,10 @@ import (
 	"go.temporal.io/sdk/workflow"
 )
 
+const (
+	CheckBeforeNotify = "checkbeforenotify"
+)
+
 type WorkflowNotifier interface {
 	Notify(workflow.Context, notifier.Info, *state.Workflow) error
 }
@@ -30,7 +35,18 @@ func (s *StateReceiver) Receive(ctx workflow.Context, c workflow.ReceiveChannel,
 }
 
 func (s *StateReceiver) Notify(ctx workflow.Context, workflowState *state.Workflow, roots map[string]RootInfo) {
-	rootInfo := roots[workflowState.ID]
+	rootInfo, ok := roots[workflowState.ID]
+
+	// TODO remove versioning
+	// there can be an edge case where we've canceled a previous check run to handle a new revision
+	// but the canceled child terraform workflow has sent over one last update prior to shutdown
+	// we should avoid notifying on this case
+	v := workflow.GetVersion(ctx, CheckBeforeNotify, workflow.DefaultVersion, workflow.Version(1))
+	if v != workflow.DefaultVersion && !ok {
+		workflow.GetLogger(ctx).Warn(fmt.Sprintf("skipping notifying root %s", workflowState.ID))
+		return
+	}
+
 	for _, notifier := range s.InternalNotifiers {
 		if err := notifier.Notify(ctx, rootInfo.ToInternalInfo(), workflowState); err != nil {
 			workflow.GetMetricsHandler(ctx).Counter("notifier_failure").Inc(1)
diff --git a/server/neptune/workflows/internal/pr/revision/state_test.go b/server/neptune/workflows/internal/pr/revision/state_test.go
index 87540b819..e4ffad9a1 100644
--- a/server/neptune/workflows/internal/pr/revision/state_test.go
+++ b/server/neptune/workflows/internal/pr/revision/state_test.go
@@ -88,7 +88,7 @@ func TestStateReceive(t *testing.T) {
 		},
 	}
 
-	t.Run("calls notifiers with state", func(t *testing.T) {
+	t.Run("calls notifiers with missing root info", func(t *testing.T) {
 		ts := testsuite.WorkflowTestSuite{}
 		env := ts.NewTestWorkflowEnvironment()
 
@@ -98,6 +98,30 @@ func TestStateReceive(t *testing.T) {
 					Output: jobOutput,
 					Status: state.WaitingJobStatus,
 				},
+				ID: internalRootInfo.ID.String(),
+			},
+			T: t,
+		})
+
+		env.AssertExpectations(t)
+
+		var result stateReceiveResponse
+		err = env.GetWorkflowResult(&result)
+		assert.False(t, result.NotifierCalled)
+		assert.NoError(t, err)
+	})
+
+	t.Run("calls notifiers with valid state", func(t *testing.T) {
+		ts := testsuite.WorkflowTestSuite{}
+		env := ts.NewTestWorkflowEnvironment()
+
+		env.ExecuteWorkflow(testStateReceiveWorkflow, stateReceiveRequest{
+			State: &state.Workflow{
+				Plan: &state.Job{
+					Output: jobOutput,
+					Status: state.WaitingJobStatus,
+				},
+				ID: internalRootInfo.ID.String(),
 			},
 			RootInfo: internalRootInfo,
 			T:        t,