mirror of
https://github.com/hibiken/asynq.git
synced 2025-10-20 09:16:12 +08:00
use sorted-set to release stale locks
This commit is contained in:
ajatprabha
87
internal/context/context.go
Normal file
87
internal/context/context.go
Normal file
@@ -0,0 +1,87 @@
|
||||
// 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 context
|
||||
|
||||
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
|
||||
qname string
|
||||
}
|
||||
|
||||
// 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
|
||||
|
||||
// New returns a context and cancel function for a given task message.
|
||||
func New(msg *base.TaskMessage, deadline time.Time) (context.Context, context.CancelFunc) {
|
||||
metadata := taskMetadata{
|
||||
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)
|
||||
}
|
||||
|
||||
// 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
|
||||
}
|
||||
|
||||
// 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
|
||||
}
|
||||
160
internal/context/context_test.go
Normal file
160
internal/context/context_test.go
Normal file
@@ -0,0 +1,160 @@
|
||||
// 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 context
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/google/go-cmp/cmp"
|
||||
"github.com/google/uuid"
|
||||
"github.com/hibiken/asynq/internal/base"
|
||||
)
|
||||
|
||||
func TestCreateContextWithFutureDeadline(t *testing.T) {
|
||||
tests := []struct {
|
||||
deadline time.Time
|
||||
}{
|
||||
{time.Now().Add(time.Hour)},
|
||||
}
|
||||
|
||||
for _, tc := range tests {
|
||||
msg := &base.TaskMessage{
|
||||
Type: "something",
|
||||
ID: uuid.New(),
|
||||
Payload: nil,
|
||||
}
|
||||
|
||||
ctx, cancel := New(msg, tc.deadline)
|
||||
|
||||
select {
|
||||
case x := <-ctx.Done():
|
||||
t.Errorf("<-ctx.Done() == %v, want nothing (it should block)", x)
|
||||
default:
|
||||
}
|
||||
|
||||
got, ok := ctx.Deadline()
|
||||
if !ok {
|
||||
t.Errorf("ctx.Deadline() returned false, want deadline to be set")
|
||||
}
|
||||
if !cmp.Equal(tc.deadline, got) {
|
||||
t.Errorf("ctx.Deadline() returned %v, want %v", got, tc.deadline)
|
||||
}
|
||||
|
||||
cancel()
|
||||
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
default:
|
||||
t.Errorf("ctx.Done() blocked, want it to be non-blocking")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestCreateContextWithPastDeadline(t *testing.T) {
|
||||
tests := []struct {
|
||||
deadline time.Time
|
||||
}{
|
||||
{time.Now().Add(-2 * time.Hour)},
|
||||
}
|
||||
|
||||
for _, tc := range tests {
|
||||
msg := &base.TaskMessage{
|
||||
Type: "something",
|
||||
ID: uuid.New(),
|
||||
Payload: nil,
|
||||
}
|
||||
|
||||
ctx, cancel := New(msg, tc.deadline)
|
||||
defer cancel()
|
||||
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
default:
|
||||
t.Errorf("ctx.Done() blocked, want it to be non-blocking")
|
||||
}
|
||||
|
||||
got, ok := ctx.Deadline()
|
||||
if !ok {
|
||||
t.Errorf("ctx.Deadline() returned false, want deadline to be set")
|
||||
}
|
||||
if !cmp.Equal(tc.deadline, got) {
|
||||
t.Errorf("ctx.Deadline() returned %v, want %v", got, tc.deadline)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetTaskMetadataFromContext(t *testing.T) {
|
||||
tests := []struct {
|
||||
desc string
|
||||
msg *base.TaskMessage
|
||||
}{
|
||||
{"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 {
|
||||
ctx, cancel := New(tc.msg, time.Now().Add(30*time.Minute))
|
||||
defer cancel()
|
||||
|
||||
id, ok := GetTaskID(ctx)
|
||||
if !ok {
|
||||
t.Errorf("%s: GetTaskID(ctx) returned ok == false", tc.desc)
|
||||
}
|
||||
if ok && id != tc.msg.ID.String() {
|
||||
t.Errorf("%s: GetTaskID(ctx) returned id == %q, want %q", tc.desc, id, tc.msg.ID.String())
|
||||
}
|
||||
|
||||
retried, ok := GetRetryCount(ctx)
|
||||
if !ok {
|
||||
t.Errorf("%s: GetRetryCount(ctx) returned ok == false", tc.desc)
|
||||
}
|
||||
if ok && retried != tc.msg.Retried {
|
||||
t.Errorf("%s: GetRetryCount(ctx) returned n == %d want %d", tc.desc, retried, tc.msg.Retried)
|
||||
}
|
||||
|
||||
maxRetry, ok := GetMaxRetry(ctx)
|
||||
if !ok {
|
||||
t.Errorf("%s: GetMaxRetry(ctx) returned ok == false", tc.desc)
|
||||
}
|
||||
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)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetTaskMetadataFromContextError(t *testing.T) {
|
||||
tests := []struct {
|
||||
desc string
|
||||
ctx context.Context
|
||||
}{
|
||||
{"with background context", context.Background()},
|
||||
}
|
||||
|
||||
for _, tc := range tests {
|
||||
if _, ok := GetTaskID(tc.ctx); ok {
|
||||
t.Errorf("%s: GetTaskID(ctx) returned ok == true", tc.desc)
|
||||
}
|
||||
if _, ok := GetRetryCount(tc.ctx); ok {
|
||||
t.Errorf("%s: GetRetryCount(ctx) returned ok == true", tc.desc)
|
||||
}
|
||||
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)
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user