Add ps command to asynqmon

2025-10-03 05:12:01 +08:00 · 2020-02-01 22:22:48 -08:00
parent d03fa34eaf
commit 96f06ac89b
15 changed files with 479 additions and 144 deletions
--- a/internal/asynqtest/asynqtest.go
+++ b/internal/asynqtest/asynqtest.go
@@ -41,6 +41,18 @@ var SortZSetEntryOpt = cmp.Transformer("SortZSetEntries", func(in []ZSetEntry) [
 	return out
 })

+// SortProcessInfoOpt is a cmp.Option to sort base.ProcessInfo for comparing slice of process info.
+var SortProcessInfoOpt = cmp.Transformer("SortProcessInfo", func(in []*base.ProcessInfo) []*base.ProcessInfo {
+	out := append([]*base.ProcessInfo(nil), in...) // Copy input to avoid mutating it
+	sort.Slice(out, func(i, j int) bool {
+		if out[i].Host != out[j].Host {
+			return out[i].Host < out[j].Host
+		}
+		return out[i].PID < out[j].PID
+	})
+	return out
+})
+
 // IgnoreIDOpt is an cmp.Option to ignore ID field in task messages when comparing.
 var IgnoreIDOpt = cmpopts.IgnoreFields(base.TaskMessage{}, "ID")

--- a/internal/base/base.go
+++ b/internal/base/base.go
@@ -8,6 +8,7 @@ package base
 import (
 	"fmt"
 	"strings"
+	"sync"
 	"time"

 	"github.com/rs/xid"
@@ -19,6 +20,7 @@ const DefaultQueueName = "default"
 // Redis keys
 const (
 	psPrefix        = "asynq:ps:"                    // HASH
+	AllProcesses    = "asynq:ps"                     // ZSET
 	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>
@@ -47,8 +49,8 @@ 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 {
+// ProcessInfoKey returns a redis key string for process info.
+func ProcessInfoKey(hostname string, pid int) string {
 	return fmt.Sprintf("%s%s:%d", psPrefix, hostname, pid)
 }

@@ -77,12 +79,47 @@ type TaskMessage struct {
 	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
+// ProcessInfo holds information about running background worker process.
+type ProcessInfo struct {
+	mu                sync.Mutex
+	Concurrency       int
+	Queues            map[string]uint
+	StrictPriority    bool
+	PID               int
+	Host              string
+	State             string
+	Started           time.Time
+	ActiveWorkerCount int
+}
+
+// NewProcessInfo returns a new instance of ProcessInfo.
+func NewProcessInfo(host string, pid, concurrency int, queues map[string]uint, strict bool) *ProcessInfo {
+	return &ProcessInfo{
+		Host:           host,
+		PID:            pid,
+		Concurrency:    concurrency,
+		Queues:         queues,
+		StrictPriority: strict,
+	}
+}
+
+// SetState set the state field of the process info.
+func (p *ProcessInfo) SetState(state string) {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	p.State = state
+}
+
+// SetStarted set the started field of the process info.
+func (p *ProcessInfo) SetStarted(t time.Time) {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	p.Started = t
+}
+
+// IncrActiveWorkerCount increments active worker count by delta.
+func (p *ProcessInfo) IncrActiveWorkerCount(delta int) {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	p.ActiveWorkerCount += delta
 }
--- a/internal/base/base_test.go
+++ b/internal/base/base_test.go
@@ -5,6 +5,7 @@
 package base

 import (
+	"sync"
 	"testing"
 	"time"
 )
@@ -61,7 +62,7 @@ func TestFailureKey(t *testing.T) {
 	}
 }

-func TestProcessStatusKey(t *testing.T) {
+func TestProcessInfoKey(t *testing.T) {
 	tests := []struct {
 		hostname string
 		pid      int
@@ -72,9 +73,36 @@ func TestProcessStatusKey(t *testing.T) {
 	}

 	for _, tc := range tests {
-		got := ProcessStatusKey(tc.hostname, tc.pid)
+		got := ProcessInfoKey(tc.hostname, tc.pid)
 		if got != tc.want {
-			t.Errorf("ProcessStatusKey(%s, %d) = %s, want %s", tc.hostname, tc.pid, got, tc.want)
+			t.Errorf("ProcessInfoKey(%s, %d) = %s, want %s", tc.hostname, tc.pid, got, tc.want)
 		}
 	}
 }
+
+// Note: Run this test with -race flag to check for data race.
+func TestProcessInfoSetter(t *testing.T) {
+	pi := NewProcessInfo("localhost", 1234, 8, map[string]uint{"default": 1}, false)
+
+	var wg sync.WaitGroup
+
+	wg.Add(3)
+
+	go func() {
+		pi.SetState("runnning")
+		wg.Done()
+	}()
+
+	go func() {
+		pi.SetStarted(time.Now())
+		pi.IncrActiveWorkerCount(1)
+		wg.Done()
+	}()
+
+	go func() {
+		pi.SetState("stopped")
+		wg.Done()
+	}()
+
+	wg.Wait()
+}
--- a/internal/rdb/inspect.go
+++ b/internal/rdb/inspect.go
@@ -755,3 +755,40 @@ func (r *RDB) RemoveQueue(qname string, force bool) error {
 	}
 	return nil
 }
+
+// ListProcesses returns the list of process statuses.
+func (r *RDB) ListProcesses() ([]*base.ProcessInfo, error) {
+	// Note: Script also removes stale keys.
+	script := redis.NewScript(`
+	local res = {}
+	local now = tonumber(ARGV[1])
+	local keys = redis.call("ZRANGEBYSCORE", KEYS[1], now, "+inf")
+	for _, key in ipairs(keys) do
+		local ps = redis.call("GET", key)
+		if ps then
+			table.insert(res, ps)
+		end  
+	end
+	redis.call("ZREMRANGEBYSCORE", KEYS[1], "-inf", now-1)
+	return res
+	`)
+	res, err := script.Run(r.client,
+		[]string{base.AllProcesses}, time.Now().UTC().Unix()).Result()
+	if err != nil {
+		return nil, err
+	}
+	data, err := cast.ToStringSliceE(res)
+	if err != nil {
+		return nil, err
+	}
+	var processes []*base.ProcessInfo
+	for _, s := range data {
+		var ps base.ProcessInfo
+		err := json.Unmarshal([]byte(s), &ps)
+		if err != nil {
+			continue // skip bad data
+		}
+		processes = append(processes, &ps)
+	}
+	return processes, nil
+}
--- a/internal/rdb/inspect_test.go
+++ b/internal/rdb/inspect_test.go
@@ -2050,3 +2050,56 @@ func TestRemoveQueueError(t *testing.T) {
 		}
 	}
 }
+
+func TestListProcesses(t *testing.T) {
+	r := setup(t)
+
+	ps1 := &base.ProcessInfo{
+		Concurrency:       10,
+		Queues:            map[string]uint{"default": 1},
+		Host:              "do.droplet1",
+		PID:               1234,
+		State:             "running",
+		Started:           time.Now().Add(-time.Hour),
+		ActiveWorkerCount: 5,
+	}
+
+	ps2 := &base.ProcessInfo{
+		Concurrency:       20,
+		Queues:            map[string]uint{"email": 1},
+		Host:              "do.droplet2",
+		PID:               9876,
+		State:             "stopped",
+		Started:           time.Now().Add(-2 * time.Hour),
+		ActiveWorkerCount: 20,
+	}
+
+	tests := []struct {
+		processes []*base.ProcessInfo
+	}{
+		{processes: []*base.ProcessInfo{}},
+		{processes: []*base.ProcessInfo{ps1}},
+		{processes: []*base.ProcessInfo{ps1, ps2}},
+	}
+
+	ignoreOpt := cmpopts.IgnoreUnexported(base.ProcessInfo{})
+
+	for _, tc := range tests {
+		h.FlushDB(t, r.client)
+
+		for _, ps := range tc.processes {
+			if err := r.WriteProcessInfo(ps, 5*time.Second); err != nil {
+				t.Fatal(err)
+			}
+		}
+
+		got, err := r.ListProcesses()
+		if err != nil {
+			t.Errorf("r.ListProcesses returned an error: %v", err)
+		}
+		if diff := cmp.Diff(tc.processes, got, h.SortProcessInfoOpt, ignoreOpt); diff != "" {
+			t.Errorf("r.ListProcesses returned %v, want %v; (-want,+got)\n%s",
+				got, tc.processes, diff)
+		}
+	}
+}
--- a/internal/rdb/rdb.go
+++ b/internal/rdb/rdb.go
@@ -347,28 +347,52 @@ func (r *RDB) forwardSingle(src, dst string) error {
 		[]string{src, dst}, now).Err()
 }

-// WriteProcessStatus writes process information to redis with expiration
+// WriteProcessInfo writes process information to redis with expiration
 // set to the value ttl.
-func (r *RDB) WriteProcessStatus(ps *base.ProcessStatus, ttl time.Duration) error {
+func (r *RDB) WriteProcessInfo(ps *base.ProcessInfo, 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()
+	// Note: Add key to ZSET with expiration time as score.
+	// ref: https://github.com/antirez/redis/issues/135#issuecomment-2361996
+	exp := time.Now().Add(ttl).UTC()
+	key := base.ProcessInfoKey(ps.Host, ps.PID)
+	// KEYS[1] -> asynq:ps
+	// KEYS[2] -> asynq:ps:<host:pid>
+	// ARGV[1] -> expiration time
+	// ARGV[2] -> TTL in seconds
+	// ARGV[3] -> process info
+	script := redis.NewScript(`
+	redis.call("ZADD", KEYS[1], ARGV[1], KEYS[2])
+	redis.call("SETEX", KEYS[2], ARGV[2], ARGV[3])
+	return redis.status_reply("OK")
+	`)
+	return script.Run(r.client, []string{base.AllProcesses, key}, float64(exp.Unix()), ttl.Seconds(), string(bytes)).Err()
 }

-// ReadProcessStatus reads process information stored in redis.
-func (r *RDB) ReadProcessStatus(host string, pid int) (*base.ProcessStatus, error) {
-	key := base.ProcessStatusKey(host, pid)
+// ReadProcessInfo reads process information stored in redis.
+func (r *RDB) ReadProcessInfo(host string, pid int) (*base.ProcessInfo, error) {
+	key := base.ProcessInfoKey(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)
+	var pinfo base.ProcessInfo
+	err = json.Unmarshal([]byte(data), &pinfo)
 	if err != nil {
 		return nil, err
 	}
-	return &ps, nil
+	return &pinfo, nil
+}
+
+// ClearProcessInfo deletes process information from redis.
+func (r *RDB) ClearProcessInfo(ps *base.ProcessInfo) error {
+	key := base.ProcessInfoKey(ps.Host, ps.PID)
+	script := redis.NewScript(`
+	redis.call("ZREM", KEYS[1], KEYS[2])
+	redis.call("DEL", KEYS[2])
+	return redis.status_reply("OK")
+	`)
+	return script.Run(r.client, []string{base.AllProcesses, key}).Err()
 }
--- a/internal/rdb/rdb_test.go
+++ b/internal/rdb/rdb_test.go
@@ -6,11 +6,13 @@ package rdb

 import (
 	"fmt"
+	"strconv"
 	"testing"
 	"time"

 	"github.com/go-redis/redis/v7"
 	"github.com/google/go-cmp/cmp"
+	"github.com/google/go-cmp/cmp/cmpopts"
 	h "github.com/hibiken/asynq/internal/asynqtest"
 	"github.com/hibiken/asynq/internal/base"
 )
@@ -739,48 +741,81 @@ func TestCheckAndEnqueue(t *testing.T) {
 	}
 }

-func TestReadWriteProcessStatus(t *testing.T) {
+func TestReadWriteClearProcessInfo(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(),
+	pinfo := &base.ProcessInfo{
+		Concurrency:       10,
+		Queues:            map[string]uint{"default": 2, "email": 5, "low": 1},
+		PID:               98765,
+		Host:              "localhost",
+		State:             "running",
+		Started:           time.Now(),
+		ActiveWorkerCount: 1,
 	}

 	tests := []struct {
-		ps  *base.ProcessStatus
+		pi  *base.ProcessInfo
 		ttl time.Duration
 	}{
-		{ps1, 5 * time.Second},
+		{pinfo, 5 * time.Second},
 	}

 	for _, tc := range tests {
 		h.FlushDB(t, r.client)

-		err := r.WriteProcessStatus(tc.ps, tc.ttl)
+		err := r.WriteProcessInfo(tc.pi, tc.ttl)
 		if err != nil {
-			t.Errorf("r.WriteProcessStatus returned an error: %v", err)
+			t.Errorf("r.WriteProcessInfo returned an error: %v", err)
 			continue
 		}

-		got, err := r.ReadProcessStatus(tc.ps.Host, tc.ps.PID)
+		got, err := r.ReadProcessInfo(tc.pi.Host, tc.pi.PID)
 		if err != nil {
-			t.Errorf("r.ReadProcessStatus returned an error: %v", err)
+			t.Errorf("r.ReadProcessInfo 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)
+		ignoreOpt := cmpopts.IgnoreUnexported(base.ProcessInfo{})
+		if diff := cmp.Diff(tc.pi, got, ignoreOpt); diff != "" {
+			t.Errorf("r.ReadProcessInfo(%q, %d) = %+v, want %+v; (-want,+got)\n%s",
+				tc.pi.Host, tc.pi.PID, got, tc.pi, diff)
 		}

-		key := base.ProcessStatusKey(tc.ps.Host, tc.ps.PID)
+		key := base.ProcessInfoKey(tc.pi.Host, tc.pi.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)
 		}
+
+		now := time.Now().UTC()
+		allKeys, err := r.client.ZRangeByScore(base.AllProcesses, &redis.ZRangeBy{
+			Min: strconv.Itoa(int(now.Unix())),
+			Max: "+inf",
+		}).Result()
+		if err != nil {
+			t.Errorf("redis ZRANGEBYSCORE %q %d +inf returned an error: %v",
+				base.AllProcesses, now.Unix(), err)
+			continue
+		}
+
+		wantAllKeys := []string{key}
+		if diff := cmp.Diff(wantAllKeys, allKeys); diff != "" {
+			t.Errorf("all keys = %v, want %v; (-want,+got)\n%s", allKeys, wantAllKeys, diff)
+		}
+
+		if err := r.ClearProcessInfo(tc.pi); err != nil {
+			t.Errorf("r.ClearProcessInfo returned an error: %v", err)
+			continue
+		}
+
+		// 1 means key exists
+		if r.client.Exists(key).Val() == 1 {
+			t.Errorf("expected %q to be deleted", key)
+		}
+
+		if r.client.ZCard(base.AllProcesses).Val() != 0 {
+			t.Errorf("expected %q to be empty", base.AllProcesses)
+		}
+
 	}
 }