diff --git a/contextutils/context.go b/contextutils/context.go index 9bc6a48d..3b8224a7 100644 --- 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 index e2effe3a..f6cb8427 100644 --- 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")