adding signal to context (#137)

Signed-off-by: Daniel Rammer <[email protected]> Signed-off-by: Daniel Rammer <[email protected]>
flyteorg · Dec 2, 2022 · 80d1624 · 80d1624
1 parent ab23187
commit 80d1624
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 0 deletions.
diff --git a/contextutils/context.go b/contextutils/context.go
@@ -24,6 +24,7 @@ const (
 	RoutineLabelKey    Key = "routine"
 	LaunchPlanIDKey    Key = "lp"
 	ResourceVersionKey Key = "res_ver"
+	SignalIDKey        Key = "signal"
 )
 
 func (k Key) String() string {
@@ -126,6 +127,11 @@ func WithTaskType(ctx context.Context, taskType string) context.Context {
 	return context.WithValue(ctx, TaskTypeKey, taskType)
 }
 
+// Gets a new context with SignalID set.
+func WithSignalID(ctx context.Context, signalID string) context.Context {
+	return context.WithValue(ctx, SignalIDKey, signalID)
+}
+
 // Gets a new context with Go Routine label key set and a label assigned to the context using pprof.Labels.
 // You can then call pprof.SetGoroutineLabels(ctx) to annotate the current go-routine and have that show up in
 // pprof analysis.

diff --git a/contextutils/context_test.go b/contextutils/context_test.go
@@ -102,6 +102,13 @@ func TestWithTaskID(t *testing.T) {
 	assert.Equal(t, "task", ctx.Value(TaskIDKey))
 }
 
+func TestWithSignalID(t *testing.T) {
+	ctx := context.Background()
+	assert.Nil(t, ctx.Value(SignalIDKey))
+	ctx = WithSignalID(ctx, "signal")
+	assert.Equal(t, "signal", ctx.Value(SignalIDKey))
+}
+
 func TestGetFields(t *testing.T) {
 	ctx := context.Background()
 	ctx = WithJobID(WithNamespace(ctx, "ns123"), "job123")