go-pkg/asynq
2
0
Fork 0
mirror of https://github.com/hibiken/asynq.git synced 2024-12-27 00:02:19 +08:00
Code Issues Projects Releases Wiki Activity

Add GetQueueName helper to extract queue name from context

Browse Source
This commit is contained in:
Ken Hibino 2020-09-07 10:54:42 -07:00
parent e09870a08a
commit 818c2d6f35
2 changed files with 27 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
context.go
View File

@ -16,6 +16,7 @@ type taskMetadata struct {
id string id string
maxRetry int maxRetry int
retryCount int retryCount int
qname string
} }
// ctxKey type is unexported to prevent collisions with context keys defined in // 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(), id: msg.ID.String(),
maxRetry: msg.Retry, maxRetry: msg.Retry,
retryCount: msg.Retried, retryCount: msg.Retried,
qname: msg.Queue,
} }
ctx := context.WithValue(context.Background(), metadataCtxKey, metadata) ctx := context.WithValue(context.Background(), metadataCtxKey, metadata)
return context.WithDeadline(ctx, deadline) return context.WithDeadline(ctx, deadline)
@ -72,3 +74,14 @@ func GetMaxRetry(ctx context.Context) (n int, ok bool) {
} }
return metadata.maxRetry, true 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
}

16
context_test.go
View File

@ -92,8 +92,9 @@ func TestGetTaskMetadataFromContext(t *testing.T) {
desc string desc string
msg *base.TaskMessage msg *base.TaskMessage
}{ }{
{"with zero retried message", &base.TaskMessage{Type: "something", ID: uuid.New(), Retry: 25, Retried: 0, 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}}, {"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 { for _, tc := range tests {
@ -123,6 +124,14 @@ func TestGetTaskMetadataFromContext(t *testing.T) {
if ok && maxRetry != tc.msg.Retry { if ok && maxRetry != tc.msg.Retry {
t.Errorf("%s: GetMaxRetry(ctx) returned n == %d want %d", tc.desc, 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 { if _, ok := GetMaxRetry(tc.ctx); ok {
t.Errorf("%s: GetMaxRetry(ctx) returned ok == true", tc.desc) 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)
}
} }
} }