github.com/wfusion/gofusion@v1.1.14/common/infra/asynq/pkg/context/context.go (about)

     1  // Copyright 2020 Kentaro Hibino. All rights reserved.
     2  // Use of this source code is governed by a MIT license
     3  // that can be found in the LICENSE file.
     4  
     5  package context
     6  
     7  import (
     8  	"context"
     9  	"time"
    10  
    11  	"github.com/wfusion/gofusion/common/infra/asynq/pkg/base"
    12  )
    13  
    14  // A taskMetadata holds task scoped data to put in context.
    15  type taskMetadata struct {
    16  	id         string
    17  	maxRetry   int
    18  	retryCount int
    19  	qname      string
    20  }
    21  
    22  // ctxKey type is unexported to prevent collisions with context keys defined in
    23  // other packages.
    24  type ctxKey int
    25  
    26  // metadataCtxKey is the context key for the task metadata.
    27  // Its value of zero is arbitrary.
    28  const metadataCtxKey ctxKey = 0
    29  
    30  // New returns a context and cancel function for a given task message.
    31  func New(base context.Context, msg *base.TaskMessage, deadline time.Time) (context.Context, context.CancelFunc) {
    32  	metadata := taskMetadata{
    33  		id:         msg.ID,
    34  		maxRetry:   msg.Retry,
    35  		retryCount: msg.Retried,
    36  		qname:      msg.Queue,
    37  	}
    38  	ctx := context.WithValue(base, metadataCtxKey, metadata)
    39  	return context.WithDeadline(ctx, deadline)
    40  }
    41  
    42  // GetTaskID extracts a task ID from a context, if any.
    43  //
    44  // ID of a task is guaranteed to be unique.
    45  // ID of a task doesn't change if the task is being retried.
    46  func GetTaskID(ctx context.Context) (id string, ok bool) {
    47  	metadata, ok := ctx.Value(metadataCtxKey).(taskMetadata)
    48  	if !ok {
    49  		return "", false
    50  	}
    51  	return metadata.id, true
    52  }
    53  
    54  // GetRetryCount extracts retry count from a context, if any.
    55  //
    56  // Return value n indicates the number of times associated task has been
    57  // retried so far.
    58  func GetRetryCount(ctx context.Context) (n int, ok bool) {
    59  	metadata, ok := ctx.Value(metadataCtxKey).(taskMetadata)
    60  	if !ok {
    61  		return 0, false
    62  	}
    63  	return metadata.retryCount, true
    64  }
    65  
    66  // GetMaxRetry extracts maximum retry from a context, if any.
    67  //
    68  // Return value n indicates the maximum number of times the associated task
    69  // can be retried if ProcessTask returns a non-nil error.
    70  func GetMaxRetry(ctx context.Context) (n int, ok bool) {
    71  	metadata, ok := ctx.Value(metadataCtxKey).(taskMetadata)
    72  	if !ok {
    73  		return 0, false
    74  	}
    75  	return metadata.maxRetry, true
    76  }
    77  
    78  // GetQueueName extracts queue name from a context, if any.
    79  //
    80  // Return value qname indicates which queue the task was pulled from.
    81  func GetQueueName(ctx context.Context) (qname string, ok bool) {
    82  	metadata, ok := ctx.Value(metadataCtxKey).(taskMetadata)
    83  	if !ok {
    84  		return "", false
    85  	}
    86  	return metadata.qname, true
    87  }