diff --git a/context.go b/context.go index 1ed6634..4710bbd 100644 --- a/context.go +++ b/context.go @@ -16,6 +16,7 @@ type taskMetadata struct { id string maxRetry int retryCount int + qname string } // ctxKey type is unexported to prevent collisions with context keys defined in @@ -32,6 +33,7 @@ func createContext(msg *base.TaskMessage, deadline time.Time) (context.Context, id: msg.ID.String(), maxRetry: msg.Retry, retryCount: msg.Retried, + qname: msg.Queue, } ctx := context.WithValue(context.Background(), metadataCtxKey, metadata) return context.WithDeadline(ctx, deadline) @@ -72,3 +74,14 @@ func GetMaxRetry(ctx context.Context) (n int, ok bool) { } return metadata.maxRetry, true } + +// GetQueueName extracts queue name from a context, if any. +// +// Return value qname indicates which queue the task was pulled from. +func GetQueueName(ctx context.Context) (qname string, ok bool) { + metadata, ok := ctx.Value(metadataCtxKey).(taskMetadata) + if !ok { + return "", false + } + return metadata.qname, true +} diff --git a/context_test.go b/context_test.go index af11d1d..305cd19 100644 --- a/context_test.go +++ b/context_test.go @@ -92,8 +92,9 @@ func TestGetTaskMetadataFromContext(t *testing.T) { desc string msg *base.TaskMessage }{ - {"with zero retried message", &base.TaskMessage{Type: "something", ID: uuid.New(), Retry: 25, Retried: 0, Timeout: 1800}}, - {"with non-zero retried message", &base.TaskMessage{Type: "something", ID: uuid.New(), Retry: 10, Retried: 5, Timeout: 1800}}, + {"with zero retried message", &base.TaskMessage{Type: "something", ID: uuid.New(), Retry: 25, Retried: 0, Timeout: 1800, Queue: "default"}}, + {"with non-zero retried message", &base.TaskMessage{Type: "something", ID: uuid.New(), Retry: 10, Retried: 5, Timeout: 1800, Queue: "default"}}, + {"with custom queue name", &base.TaskMessage{Type: "something", ID: uuid.New(), Retry: 25, Retried: 0, Timeout: 1800, Queue: "custom"}}, } for _, tc := range tests { @@ -123,6 +124,14 @@ func TestGetTaskMetadataFromContext(t *testing.T) { if ok && maxRetry != tc.msg.Retry { t.Errorf("%s: GetMaxRetry(ctx) returned n == %d want %d", tc.desc, maxRetry, tc.msg.Retry) } + + qname, ok := GetQueueName(ctx) + if !ok { + t.Errorf("%s: GetQueueName(ctx) returned ok == false", tc.desc) + } + if ok && qname != tc.msg.Queue { + t.Errorf("%s: GetQueueName(ctx) returned qname == %q, want %q", tc.desc, qname, tc.msg.Queue) + } } } @@ -144,5 +153,8 @@ func TestGetTaskMetadataFromContextError(t *testing.T) { if _, ok := GetMaxRetry(tc.ctx); ok { t.Errorf("%s: GetMaxRetry(ctx) returned ok == true", tc.desc) } + if _, ok := GetQueueName(tc.ctx); ok { + t.Errorf("%s: GetQueueName(ctx) returned ok == true", tc.desc) + } } }