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

Refactor redis keys and store messages in protobuf

Browse Source 
Changes:
- Task messages are stored under "asynq:{<qname>}:t:<task_id>" key in redis, value is a HASH type and message are stored under "msg" key in the hash. The hash also stores "deadline", "timeout".
- Redis LIST and ZSET stores task message IDs
- Task messages are serialized using protocol buffer
This commit is contained in:
Ken Hibino
2021-03-12 16:23:08 -08:00
parent 2516c4baba
commit 7af3981929
22 changed files with 2534 additions and 668 deletions
Download Patch File Download Diff File Expand all files Collapse all files

297
internal/base/base.go
View File

@@ -6,6 +6,7 @@
package base
import (
"bytes"
"context"
"encoding/json"
"fmt"
@@ -15,7 +16,10 @@ import (
"time"
"github.com/go-redis/redis/v7"
"github.com/golang/protobuf/ptypes"
"github.com/google/uuid"
pb "github.com/hibiken/asynq/internal/proto"
"google.golang.org/protobuf/proto"
)
// Version of asynq library and CLI.
@@ -25,7 +29,7 @@ const Version = "0.17.2"
const DefaultQueueName = "default"
// DefaultQueue is the redis key for the default queue.
var DefaultQueue = QueueKey(DefaultQueueName)
var DefaultQueue = PendingKey(DefaultQueueName)
// Global Redis keys.
const (
@@ -45,9 +49,19 @@ func ValidateQueueName(qname string) error {
return nil
}
// QueueKey returns a redis key for the given queue name.
func QueueKey(qname string) string {
return fmt.Sprintf("asynq:{%s}", qname)
// TaskKeyPrefix returns a prefix for task key.
func TaskKeyPrefix(qname string) string {
return fmt.Sprintf("asynq:{%s}:t:", qname)
}
// TaskKey returns a redis key for the given task message.
func TaskKey(qname, id string) string {
return fmt.Sprintf("%s%s", TaskKeyPrefix(qname), id)
}
// PendingKey returns a redis key for the given queue name.
func PendingKey(qname string) string {
return fmt.Sprintf("asynq:{%s}:pending", qname)
}
// ActiveKey returns a redis key for the active tasks.
@@ -184,24 +198,51 @@ type TaskMessage struct {
UniqueKey string
}
// EncodeMessage marshals the given task message in JSON and returns an encoded string.
func EncodeMessage(msg *TaskMessage) (string, error) {
b, err := json.Marshal(msg)
if err != nil {
return "", err
// EncodeMessage marshals the given task message and returns an encoded bytes.
func EncodeMessage(msg *TaskMessage) ([]byte, error) {
if msg == nil {
return nil, fmt.Errorf("cannot encode nil message")
}
return string(b), nil
}
// DecodeMessage unmarshals the given encoded string and returns a decoded task message.
func DecodeMessage(s string) (*TaskMessage, error) {
d := json.NewDecoder(strings.NewReader(s))
d.UseNumber()
var msg TaskMessage
if err := d.Decode(&msg); err != nil {
payload, err := json.Marshal(msg.Payload)
if err != nil {
return nil, err
}
return &msg, nil
return proto.Marshal(&pb.TaskMessage{
Type: msg.Type,
Payload: payload,
Id: msg.ID.String(),
Queue: msg.Queue,
Retry: int32(msg.Retry),
Retried: int32(msg.Retried),
ErrorMsg: msg.ErrorMsg,
Timeout: msg.Timeout,
Deadline: msg.Deadline,
UniqueKey: msg.UniqueKey,
})
}
// DecodeMessage unmarshals the given bytes and returns a decoded task message.
func DecodeMessage(data []byte) (*TaskMessage, error) {
var pbmsg pb.TaskMessage
if err := proto.Unmarshal(data, &pbmsg); err != nil {
return nil, err
}
payload, err := decodePayload(pbmsg.GetPayload())
if err != nil {
return nil, err
}
return &TaskMessage{
Type: pbmsg.GetType(),
Payload: payload,
ID: uuid.MustParse(pbmsg.GetId()),
Queue: pbmsg.GetQueue(),
Retry: int(pbmsg.GetRetry()),
Retried: int(pbmsg.GetRetried()),
ErrorMsg: pbmsg.GetErrorMsg(),
Timeout: pbmsg.GetTimeout(),
Deadline: pbmsg.GetDeadline(),
UniqueKey: pbmsg.GetUniqueKey(),
}, nil
}
// Z represents sorted set member.
@@ -282,6 +323,59 @@ type ServerInfo struct {
ActiveWorkerCount int
}
// EncodeServerInfo marshals the given ServerInfo and returns the encoded bytes.
func EncodeServerInfo(info *ServerInfo) ([]byte, error) {
if info == nil {
return nil, fmt.Errorf("cannot encode nil server info")
}
queues := make(map[string]int32)
for q, p := range info.Queues {
queues[q] = int32(p)
}
started, err := ptypes.TimestampProto(info.Started)
if err != nil {
return nil, err
}
return proto.Marshal(&pb.ServerInfo{
Host: info.Host,
Pid: int32(info.PID),
ServerId: info.ServerID,
Concurrency: int32(info.Concurrency),
Queues: queues,
StrictPriority: info.StrictPriority,
Status: info.Status,
StartTime: started,
ActiveWorkerCount: int32(info.ActiveWorkerCount),
})
}
// DecodeServerInfo decodes the given bytes into ServerInfo.
func DecodeServerInfo(b []byte) (*ServerInfo, error) {
var pbmsg pb.ServerInfo
if err := proto.Unmarshal(b, &pbmsg); err != nil {
return nil, err
}
queues := make(map[string]int)
for q, p := range pbmsg.GetQueues() {
queues[q] = int(p)
}
startTime, err := ptypes.Timestamp(pbmsg.GetStartTime())
if err != nil {
return nil, err
}
return &ServerInfo{
Host: pbmsg.GetHost(),
PID: int(pbmsg.GetPid()),
ServerID: pbmsg.GetServerId(),
Concurrency: int(pbmsg.GetConcurrency()),
Queues: queues,
StrictPriority: pbmsg.GetStrictPriority(),
Status: pbmsg.GetStatus(),
Started: startTime,
ActiveWorkerCount: int(pbmsg.GetActiveWorkerCount()),
}, nil
}
// WorkerInfo holds information about a running worker.
type WorkerInfo struct {
Host string
@@ -289,12 +383,83 @@ type WorkerInfo struct {
ServerID string
ID string
Type string
Queue string
Payload map[string]interface{}
Queue string
Started time.Time
Deadline time.Time
}
// EncodeWorkerInfo marshals the given WorkerInfo and returns the encoded bytes.
func EncodeWorkerInfo(info *WorkerInfo) ([]byte, error) {
if info == nil {
return nil, fmt.Errorf("cannot encode nil worker info")
}
payload, err := json.Marshal(info.Payload)
if err != nil {
return nil, err
}
startTime, err := ptypes.TimestampProto(info.Started)
if err != nil {
return nil, err
}
deadline, err := ptypes.TimestampProto(info.Deadline)
if err != nil {
return nil, err
}
return proto.Marshal(&pb.WorkerInfo{
Host: info.Host,
Pid: int32(info.PID),
ServerId: info.ServerID,
TaskId: info.ID,
TaskType: info.Type,
TaskPayload: payload,
Queue: info.Queue,
StartTime: startTime,
Deadline: deadline,
})
}
func decodePayload(b []byte) (map[string]interface{}, error) {
d := json.NewDecoder(bytes.NewReader(b))
d.UseNumber()
payload := make(map[string]interface{})
if err := d.Decode(&payload); err != nil {
return nil, err
}
return payload, nil
}
// DecodeWorkerInfo decodes the given bytes into WorkerInfo.
func DecodeWorkerInfo(b []byte) (*WorkerInfo, error) {
var pbmsg pb.WorkerInfo
if err := proto.Unmarshal(b, &pbmsg); err != nil {
return nil, err
}
payload, err := decodePayload(pbmsg.GetTaskPayload())
if err != nil {
return nil, err
}
startTime, err := ptypes.Timestamp(pbmsg.GetStartTime())
if err != nil {
return nil, err
}
deadline, err := ptypes.Timestamp(pbmsg.GetDeadline())
if err != nil {
return nil, err
}
return &WorkerInfo{
Host: pbmsg.GetHost(),
PID: int(pbmsg.GetPid()),
ServerID: pbmsg.GetServerId(),
ID: pbmsg.GetTaskId(),
Type: pbmsg.GetTaskType(),
Payload: payload,
Queue: pbmsg.GetQueue(),
Started: startTime,
Deadline: deadline,
}, nil
}
// SchedulerEntry holds information about a periodic task registered with a scheduler.
type SchedulerEntry struct {
// Identifier of this entry.
@@ -320,6 +485,63 @@ type SchedulerEntry struct {
Prev time.Time
}
// EncodeSchedulerEntry marshals the given entry and returns an encoded bytes.
func EncodeSchedulerEntry(entry *SchedulerEntry) ([]byte, error) {
if entry == nil {
return nil, fmt.Errorf("cannot encode nil scheduler entry")
}
payload, err := json.Marshal(entry.Payload)
if err != nil {
return nil, err
}
next, err := ptypes.TimestampProto(entry.Next)
if err != nil {
return nil, err
}
prev, err := ptypes.TimestampProto(entry.Prev)
if err != nil {
return nil, err
}
return proto.Marshal(&pb.SchedulerEntry{
Id: entry.ID,
Spec: entry.Spec,
TaskType: entry.Type,
TaskPayload: payload,
EnqueueOptions: entry.Opts,
NextEnqueueTime: next,
PrevEnqueueTime: prev,
})
}
// DecodeSchedulerEntry unmarshals the given bytes and returns a decoded SchedulerEntry.
func DecodeSchedulerEntry(b []byte) (*SchedulerEntry, error) {
var pbmsg pb.SchedulerEntry
if err := proto.Unmarshal(b, &pbmsg); err != nil {
return nil, err
}
payload, err := decodePayload(pbmsg.GetTaskPayload())
if err != nil {
return nil, err
}
next, err := ptypes.Timestamp(pbmsg.GetNextEnqueueTime())
if err != nil {
return nil, err
}
prev, err := ptypes.Timestamp(pbmsg.GetPrevEnqueueTime())
if err != nil {
return nil, err
}
return &SchedulerEntry{
ID: pbmsg.GetId(),
Spec: pbmsg.GetSpec(),
Type: pbmsg.GetTaskType(),
Payload: payload,
Opts: pbmsg.GetEnqueueOptions(),
Next: next,
Prev: prev,
}, nil
}
// SchedulerEnqueueEvent holds information about an enqueue event by a scheduler.
type SchedulerEnqueueEvent struct {
// ID of the task that was enqueued.
@@ -329,6 +551,39 @@ type SchedulerEnqueueEvent struct {
EnqueuedAt time.Time
}
// EncodeSchedulerEnqueueEvent marshals the given event
// and returns an encoded bytes.
func EncodeSchedulerEnqueueEvent(event *SchedulerEnqueueEvent) ([]byte, error) {
if event == nil {
return nil, fmt.Errorf("cannot encode nil enqueue event")
}
enqueuedAt, err := ptypes.TimestampProto(event.EnqueuedAt)
if err != nil {
return nil, err
}
return proto.Marshal(&pb.SchedulerEnqueueEvent{
TaskId: event.TaskID,
EnqueueTime: enqueuedAt,
})
}
// DecodeSchedulerEnqueueEvent unmarshals the given bytes
// and returns a decoded SchedulerEnqueueEvent.
func DecodeSchedulerEnqueueEvent(b []byte) (*SchedulerEnqueueEvent, error) {
var pbmsg pb.SchedulerEnqueueEvent
if err := proto.Unmarshal(b, &pbmsg); err != nil {
return nil, err
}
enqueuedAt, err := ptypes.Timestamp(pbmsg.GetEnqueueTime())
if err != nil {
return nil, err
}
return &SchedulerEnqueueEvent{
TaskID: pbmsg.GetTaskId(),
EnqueuedAt: enqueuedAt,
}, nil
}
// Cancelations is a collection that holds cancel functions for all active tasks.
//
// Cancelations are safe for concurrent use by multipel goroutines.
@@ -380,7 +635,7 @@ type Broker interface {
ScheduleUnique(msg *TaskMessage, processAt time.Time, ttl time.Duration) error
Retry(msg *TaskMessage, processAt time.Time, errMsg string) error
Archive(msg *TaskMessage, errMsg string) error
CheckAndEnqueue(qnames ...string) error
ForwardIfReady(qnames ...string) error
ListDeadlineExceeded(deadline time.Time, qnames ...string) ([]*TaskMessage, error)
WriteServerState(info *ServerInfo, workers []*WorkerInfo, ttl time.Duration) error
ClearServerState(host string, pid int, serverID string) error

165
internal/base/base_test.go
View File

@@ -7,6 +7,7 @@ package base
import (
"context"
"encoding/json"
"fmt"
"sync"
"testing"
"time"
@@ -15,17 +16,36 @@ import (
"github.com/google/uuid"
)
func TestTaskKey(t *testing.T) {
id := uuid.NewString()
tests := []struct {
qname string
id string
want string
}{
{"default", id, fmt.Sprintf("asynq:{default}:t:%s", id)},
}
for _, tc := range tests {
got := TaskKey(tc.qname, tc.id)
if got != tc.want {
t.Errorf("TaskKey(%q, %s) = %q, want %q", tc.qname, tc.id, got, tc.want)
}
}
}
func TestQueueKey(t *testing.T) {
tests := []struct {
qname string
want string
}{
{"default", "asynq:{default}"},
{"custom", "asynq:{custom}"},
{"default", "asynq:{default}:pending"},
{"custom", "asynq:{custom}:pending"},
}
for _, tc := range tests {
got := QueueKey(tc.qname)
got := PendingKey(tc.qname)
if got != tc.want {
t.Errorf("QueueKey(%q) = %q, want %q", tc.qname, got, tc.want)
}
@@ -352,6 +372,145 @@ func TestMessageEncoding(t *testing.T) {
}
}
func TestServerInfoEncoding(t *testing.T) {
tests := []struct {
info ServerInfo
}{
{
info: ServerInfo{
Host: "127.0.0.1",
PID: 9876,
ServerID: "abc123",
Concurrency: 10,
Queues: map[string]int{"default": 1, "critical": 2},
StrictPriority: false,
Status: "running",
Started: time.Now().Add(-3 * time.Hour),
ActiveWorkerCount: 8,
},
},
}
for _, tc := range tests {
encoded, err := EncodeServerInfo(&tc.info)
if err != nil {
t.Errorf("EncodeServerInfo(info) returned error: %v", err)
continue
}
decoded, err := DecodeServerInfo(encoded)
if err != nil {
t.Errorf("DecodeServerInfo(encoded) returned error: %v", err)
continue
}
if diff := cmp.Diff(&tc.info, decoded); diff != "" {
t.Errorf("Decoded ServerInfo == %+v, want %+v;(-want,+got)\n%s",
decoded, tc.info, diff)
}
}
}
func TestWorkerInfoEncoding(t *testing.T) {
tests := []struct {
info WorkerInfo
}{
{
info: WorkerInfo{
Host: "127.0.0.1",
PID: 9876,
ServerID: "abc123",
ID: uuid.NewString(),
Type: "taskA",
Payload: map[string]interface{}{"foo": "bar"},
Queue: "default",
Started: time.Now().Add(-3 * time.Hour),
Deadline: time.Now().Add(30 * time.Second),
},
},
}
for _, tc := range tests {
encoded, err := EncodeWorkerInfo(&tc.info)
if err != nil {
t.Errorf("EncodeWorkerInfo(info) returned error: %v", err)
continue
}
decoded, err := DecodeWorkerInfo(encoded)
if err != nil {
t.Errorf("DecodeWorkerInfo(encoded) returned error: %v", err)
continue
}
if diff := cmp.Diff(&tc.info, decoded); diff != "" {
t.Errorf("Decoded WorkerInfo == %+v, want %+v;(-want,+got)\n%s",
decoded, tc.info, diff)
}
}
}
func TestSchedulerEntryEncoding(t *testing.T) {
tests := []struct {
entry SchedulerEntry
}{
{
entry: SchedulerEntry{
ID: uuid.NewString(),
Spec: "* * * * *",
Type: "task_A",
Payload: map[string]interface{}{"foo": "bar"},
Opts: []string{"Queue('email')"},
Next: time.Now().Add(30 * time.Second).UTC(),
Prev: time.Now().Add(-2 * time.Minute).UTC(),
},
},
}
for _, tc := range tests {
encoded, err := EncodeSchedulerEntry(&tc.entry)
if err != nil {
t.Errorf("EncodeSchedulerEntry(entry) returned error: %v", err)
continue
}
decoded, err := DecodeSchedulerEntry(encoded)
if err != nil {
t.Errorf("DecodeSchedulerEntry(encoded) returned error: %v", err)
continue
}
if diff := cmp.Diff(&tc.entry, decoded); diff != "" {
t.Errorf("Decoded SchedulerEntry == %+v, want %+v;(-want,+got)\n%s",
decoded, tc.entry, diff)
}
}
}
func TestSchedulerEnqueueEventEncoding(t *testing.T) {
tests := []struct {
event SchedulerEnqueueEvent
}{
{
event: SchedulerEnqueueEvent{
TaskID: uuid.NewString(),
EnqueuedAt: time.Now().Add(-30 * time.Second).UTC(),
},
},
}
for _, tc := range tests {
encoded, err := EncodeSchedulerEnqueueEvent(&tc.event)
if err != nil {
t.Errorf("EncodeSchedulerEnqueueEvent(event) returned error: %v", err)
continue
}
decoded, err := DecodeSchedulerEnqueueEvent(encoded)
if err != nil {
t.Errorf("DecodeSchedulerEnqueueEvent(encoded) returned error: %v", err)
continue
}
if diff := cmp.Diff(&tc.event, decoded); diff != "" {
t.Errorf("Decoded SchedulerEnqueueEvent == %+v, want %+v;(-want,+got)\n%s",
decoded, tc.event, diff)
}
}
}
// Test for status being accessed by multiple goroutines.
// Run with -race flag to check for data race.
func TestStatusConcurrentAccess(t *testing.T) {