// Copyright 2020 Kentaro Hibino. All rights reserved. // Use of this source code is governed by a MIT license // that can be found in the LICENSE file. package asynq import ( "context" "time" "github.com/hibiken/asynq/internal/base" ) // A taskMetadata holds task scoped data to put in context. type taskMetadata struct { id string maxRetry int retryCount int } // ctxKey type is unexported to prevent collisions with context keys defined in // other packages. type ctxKey int // metadataCtxKey is the context key for the task metadata. // Its value of zero is arbitrary. const metadataCtxKey ctxKey = 0 // createContext returns a context and cancel function for a given task message. func createContext(msg *base.TaskMessage) (ctx context.Context, cancel context.CancelFunc) { metadata := taskMetadata{ id: msg.ID.String(), maxRetry: msg.Retry, retryCount: msg.Retried, } ctx = context.WithValue(context.Background(), metadataCtxKey, metadata) if msg.Timeout == 0 && msg.Deadline == 0 { panic("asynq: internal error: missing both timeout and deadline") } if msg.Timeout != 0 { timeout := time.Duration(msg.Timeout) * time.Second ctx, cancel = context.WithTimeout(ctx, timeout) } if msg.Deadline != 0 { deadline := time.Unix(int64(msg.Deadline), 0) ctx, cancel = context.WithDeadline(ctx, deadline) } if cancel == nil { ctx, cancel = context.WithCancel(ctx) } return ctx, cancel } // GetTaskID extracts a task ID from a context, if any. // // ID of a task is guaranteed to be unique. // ID of a task doesn't change if the task is being retried. func GetTaskID(ctx context.Context) (id string, ok bool) { metadata, ok := ctx.Value(metadataCtxKey).(taskMetadata) if !ok { return "", false } return metadata.id, true } // GetRetryCount extracts retry count from a context, if any. // // Return value n indicates the number of times associated task has been // retried so far. func GetRetryCount(ctx context.Context) (n int, ok bool) { metadata, ok := ctx.Value(metadataCtxKey).(taskMetadata) if !ok { return 0, false } return metadata.retryCount, true } // GetMaxRetry extracts maximum retry from a context, if any. // // Return value n indicates the maximum number of times the assoicated task // can be retried if ProcessTask returns a non-nil error. func GetMaxRetry(ctx context.Context) (n int, ok bool) { metadata, ok := ctx.Value(metadataCtxKey).(taskMetadata) if !ok { return 0, false } return metadata.maxRetry, true }