Add hearbeater

internal/base/base.go

@@ -6,6 +6,7 @@
 package base
 
 import (
+	"fmt"
 	"strings"
 	"time"
 
@@ -17,6 +18,7 @@ const DefaultQueueName = "default"
 
 // Redis keys
 const (
+	psPrefix        = "asynq:ps:"                    // HASH
 	processedPrefix = "asynq:processed:"             // STRING - asynq:processed:<yyyy-mm-dd>
 	failurePrefix   = "asynq:failure:"               // STRING - asynq:failure:<yyyy-mm-dd>
 	QueuePrefix     = "asynq:queues:"                // LIST   - asynq:queues:<qname>
@@ -45,6 +47,11 @@ func FailureKey(t time.Time) string {
 	return failurePrefix + t.UTC().Format("2006-01-02")
 }
 
+// ProcessStatusKey returns a redis key string for process status.
+func ProcessStatusKey(hostname string, pid int) string {
+	return fmt.Sprintf("%s%s:%d", psPrefix, hostname, pid)
+}
+
 // TaskMessage is the internal representation of a task with additional metadata fields.
 // Serialized data of this type gets written to redis.
 type TaskMessage struct {
@@ -69,3 +76,13 @@ type TaskMessage struct {
 	// ErrorMsg holds the error message from the last failure.
 	ErrorMsg string
 }
+
+// ProcessStatus holds information about running background worker process.
+type ProcessStatus struct {
+	Concurrency int
+	Queues      map[string]uint
+	PID         int
+	Host        string
+	State       string
+	Started     time.Time
+}

internal/base/base_test.go

@@ -60,3 +60,21 @@ func TestFailureKey(t *testing.T) {
 		}
 	}
 }
+
+func TestProcessStatusKey(t *testing.T) {
+	tests := []struct {
+		hostname string
+		pid      int
+		want     string
+	}{
+		{"localhost", 9876, "asynq:ps:localhost:9876"},
+		{"127.0.0.1", 1234, "asynq:ps:127.0.0.1:1234"},
+	}
+
+	for _, tc := range tests {
+		got := ProcessStatusKey(tc.hostname, tc.pid)
+		if got != tc.want {
+			t.Errorf("ProcessStatusKey(%s, %d) = %s, want %s", tc.hostname, tc.pid, got, tc.want)
+		}
+	}
+}

internal/rdb/rdb.go

@@ -346,3 +346,29 @@ func (r *RDB) forwardSingle(src, dst string) error {
 	return script.Run(r.client,
 		[]string{src, dst}, now).Err()
 }
+
+// WriteProcessStatus writes process information to redis with expiration
+// set to the value ttl.
+func (r *RDB) WriteProcessStatus(ps *base.ProcessStatus, ttl time.Duration) error {
+	bytes, err := json.Marshal(ps)
+	if err != nil {
+		return err
+	}
+	key := base.ProcessStatusKey(ps.Host, ps.PID)
+	return r.client.Set(key, string(bytes), ttl).Err()
+}
+
+// ReadProcessStatus reads process information stored in redis.
+func (r *RDB) ReadProcessStatus(host string, pid int) (*base.ProcessStatus, error) {
+	key := base.ProcessStatusKey(host, pid)
+	data, err := r.client.Get(key).Result()
+	if err != nil {
+		return nil, err
+	}
+	var ps base.ProcessStatus
+	err = json.Unmarshal([]byte(data), &ps)
+	if err != nil {
+		return nil, err
+	}
+	return &ps, nil
+}

internal/rdb/rdb_test.go

@@ -738,3 +738,49 @@ func TestCheckAndEnqueue(t *testing.T) {
 		}
 	}
 }
+
+func TestReadWriteProcessStatus(t *testing.T) {
+	r := setup(t)
+	ps1 := &base.ProcessStatus{
+		Concurrency: 10,
+		Queues:      map[string]uint{"default": 2, "email": 5, "low": 1},
+		PID:         98765,
+		Host:        "localhost",
+		State:       "running",
+		Started:     time.Now(),
+	}
+
+	tests := []struct {
+		ps  *base.ProcessStatus
+		ttl time.Duration
+	}{
+		{ps1, 5 * time.Second},
+	}
+
+	for _, tc := range tests {
+		h.FlushDB(t, r.client)
+
+		err := r.WriteProcessStatus(tc.ps, tc.ttl)
+		if err != nil {
+			t.Errorf("r.WriteProcessStatus returned an error: %v", err)
+			continue
+		}
+
+		got, err := r.ReadProcessStatus(tc.ps.Host, tc.ps.PID)
+		if err != nil {
+			t.Errorf("r.ReadProcessStatus returned an error: %v", err)
+			continue
+		}
+
+		if diff := cmp.Diff(tc.ps, got); diff != "" {
+			t.Errorf("r.ReadProcessStatus(%q, %d) = %+v, want %+v; (-want,+got)\n%s",
+				tc.ps.Host, tc.ps.PID, got, tc.ps, diff)
+		}
+
+		key := base.ProcessStatusKey(tc.ps.Host, tc.ps.PID)
+		gotTTL := r.client.TTL(key).Val()
+		if !cmp.Equal(tc.ttl, gotTTL, timeCmpOpt) {
+			t.Errorf("redis TTL %q returned %v, want %v", key, gotTTL, tc.ttl)
+		}
+	}
+}