Repository: hibiken/asynq
Pull request #876 (2024-05-01 00:11:42)

fix  testing errors caused by possible inconsistent scores between ProcessIn and ProcessAt


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 69414c3..c167e8d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+## [0.25.0] - 2023-01-02
+
+### Added
+- Added configuration for Janitor's Interval and Deletion Batch Size (PR: https://github.com/hibiken/asynq/pull/715)
+
 ## [0.24.1] - 2023-05-01
 
 ### Changed
diff --git a/client.go b/client.go
index 8f2f92c..aec50b1 100644
--- a/client.go
+++ b/client.go
@@ -10,11 +10,11 @@ import (
 	"strings"
 	"time"
 
-	"github.com/redis/go-redis/v9"
 	"github.com/google/uuid"
 	"github.com/hibiken/asynq/internal/base"
 	"github.com/hibiken/asynq/internal/errors"
 	"github.com/hibiken/asynq/internal/rdb"
+	"github.com/redis/go-redis/v9"
 )
 
 // A Client is responsible for scheduling tasks.
@@ -25,15 +25,26 @@ import (
 // Clients are safe for concurrent use by multiple goroutines.
 type Client struct {
 	broker base.Broker
+	// When a Client has been created with an existing Redis connection, we do
+	// not want to close it.
+	sharedConnection bool
 }
 
 // NewClient returns a new Client instance given a redis connection option.
 func NewClient(r RedisConnOpt) *Client {
-	c, ok := r.MakeRedisClient().(redis.UniversalClient)
+	redisClient, ok := r.MakeRedisClient().(redis.UniversalClient)
 	if !ok {
 		panic(fmt.Sprintf("asynq: unsupported RedisConnOpt type %T", r))
 	}
-	return &Client{broker: rdb.NewRDB(c)}
+	client := NewClientFromRedisClient(redisClient)
+	client.sharedConnection = false
+	return client
+}
+
+// NewClientFromRedisClient returns a new instance of Client given a redis.UniversalClient
+// Warning: The underlying redis connection pool will not be closed by Asynq, you are responsible for closing it.
+func NewClientFromRedisClient(c redis.UniversalClient) *Client {
+	return &Client{broker: rdb.NewRDB(c), sharedConnection: true}
 }
 
 type OptionType int
@@ -150,9 +161,9 @@ func (t deadlineOption) Value() interface{} { return time.Time(t) }
 // TTL duration must be greater than or equal to 1 second.
 //
 // Uniqueness of a task is based on the following properties:
-//     - Task Type
-//     - Task Payload
-//     - Queue Name
+//   - Task Type
+//   - Task Payload
+//   - Queue Name
 func Unique(ttl time.Duration) Option {
 	return uniqueOption(ttl)
 }
@@ -307,6 +318,9 @@ var (
 
 // Close closes the connection with redis.
 func (c *Client) Close() error {
+	if c.sharedConnection {
+		return fmt.Errorf("redis connection is shared so the Client can't be closed through asynq")
+	}
 	return c.broker.Close()
 }
 
@@ -405,6 +419,11 @@ func (c *Client) EnqueueContext(ctx context.Context, task *Task, opts ...Option)
 	return newTaskInfo(msg, state, opt.processAt, nil), nil
 }
 
+// Ping performs a ping against the redis connection.
+func (c *Client) Ping() error {
+	return c.broker.Ping()
+}
+
 func (c *Client) enqueue(ctx context.Context, msg *base.TaskMessage, uniqueTTL time.Duration) error {
 	if uniqueTTL > 0 {
 		return c.broker.EnqueueUnique(ctx, msg, uniqueTTL)
diff --git a/inspector.go b/inspector.go
index ee99f52..e361d22 100644
--- a/inspector.go
+++ b/inspector.go
@@ -10,16 +10,19 @@ import (
 	"strings"
 	"time"
 
-	"github.com/redis/go-redis/v9"
 	"github.com/hibiken/asynq/internal/base"
 	"github.com/hibiken/asynq/internal/errors"
 	"github.com/hibiken/asynq/internal/rdb"
+	"github.com/redis/go-redis/v9"
 )
 
 // Inspector is a client interface to inspect and mutate the state of
 // queues and tasks.
 type Inspector struct {
 	rdb *rdb.RDB
+	// When an Inspector has been created with an existing Redis connection, we do
+	// not want to close it.
+	sharedConnection bool
 }
 
 // New returns a new instance of Inspector.
@@ -28,13 +31,25 @@ func NewInspector(r RedisConnOpt) *Inspector {
 	if !ok {
 		panic(fmt.Sprintf("inspeq: unsupported RedisConnOpt type %T", r))
 	}
+	inspector := NewInspectorFromRedisClient(c)
+	inspector.sharedConnection = false
+	return inspector
+}
+
+// NewInspectorFromRedisClient returns a new instance of Inspector given a redis.UniversalClient
+// Warning: The underlying redis connection pool will not be closed by Asynq, you are responsible for closing it.
+func NewInspectorFromRedisClient(c redis.UniversalClient) *Inspector {
 	return &Inspector{
-		rdb: rdb.NewRDB(c),
+		rdb:              rdb.NewRDB(c),
+		sharedConnection: true,
 	}
 }
 
 // Close closes the connection with redis.
 func (i *Inspector) Close() error {
+	if i.sharedConnection {
+		return fmt.Errorf("redis connection is shared so the Inspector can't be closed through asynq")
+	}
 	return i.rdb.Close()
 }
 
diff --git a/internal/base/base.go b/internal/base/base.go
index a63c548..c0eb5bb 100644
--- a/internal/base/base.go
+++ b/internal/base/base.go
@@ -737,7 +737,7 @@ type Broker interface {
 	ReclaimStaleAggregationSets(qname string) error
 
 	// Task retention related method
-	DeleteExpiredCompletedTasks(qname string) error
+	DeleteExpiredCompletedTasks(qname string, batchSize int) error
 
 	// Lease related methods
 	ListLeaseExpired(cutoff time.Time, qnames ...string) ([]*TaskMessage, error)
diff --git a/internal/rdb/rdb.go b/internal/rdb/rdb.go
index 68de386..74c140e 100644
--- a/internal/rdb/rdb.go
+++ b/internal/rdb/rdb.go
@@ -829,6 +829,7 @@ const (
 // KEYS[6] -> asynq:{<qname>}:failed:<yyyy-mm-dd>
 // KEYS[7] -> asynq:{<qname>}:processed
 // KEYS[8] -> asynq:{<qname>}:failed
+// KEYS[9] -> asynq:{<qname>}:t:
 // -------
 // ARGV[1] -> task ID
 // ARGV[2] -> updated base.TaskMessage value
@@ -845,8 +846,22 @@ if redis.call("ZREM", KEYS[3], ARGV[1]) == 0 then
   return redis.error_reply("NOT FOUND")
 end
 redis.call("ZADD", KEYS[4], ARGV[3], ARGV[1])
-redis.call("ZREMRANGEBYSCORE", KEYS[4], "-inf", ARGV[4])
-redis.call("ZREMRANGEBYRANK", KEYS[4], 0, -ARGV[5])
+local old = redis.call("ZRANGE", KEYS[4], "-inf", ARGV[4], "BYSCORE")
+if #old > 0 then
+	for _, id in ipairs(old) do
+		redis.call("DEL", KEYS[9] .. id)
+	end
+	redis.call("ZREM", KEYS[4], unpack(old))
+end
+
+local extra = redis.call("ZRANGE", KEYS[4], 0, -ARGV[5])
+if #extra > 0 then
+	for _, id in ipairs(extra) do
+		redis.call("DEL", KEYS[9] .. id)
+	end
+	redis.call("ZREM", KEYS[4], unpack(extra))
+end
+
 redis.call("HSET", KEYS[1], "msg", ARGV[2], "state", "archived")
 local n = redis.call("INCR", KEYS[5])
 if tonumber(n) == 1 then
@@ -889,6 +904,7 @@ func (r *RDB) Archive(ctx context.Context, msg *base.TaskMessage, errMsg string)
 		base.FailedKey(msg.Queue, now),
 		base.ProcessedTotalKey(msg.Queue),
 		base.FailedTotalKey(msg.Queue),
+		base.TaskKeyPrefix(msg.Queue),
 	}
 	argv := []interface{}{
 		msg.ID,
@@ -1241,9 +1257,7 @@ return table.getn(ids)`)
 
 // DeleteExpiredCompletedTasks checks for any expired tasks in the given queue's completed set,
 // and delete all expired tasks.
-func (r *RDB) DeleteExpiredCompletedTasks(qname string) error {
-	// Note: Do this operation in fix batches to prevent long running script.
-	const batchSize = 100
+func (r *RDB) DeleteExpiredCompletedTasks(qname string, batchSize int) error {
 	for {
 		n, err := r.deleteExpiredCompletedTasks(qname, batchSize)
 		if err != nil {
diff --git a/janitor.go b/janitor.go
index 40f5ea1..612dbf0 100644
--- a/janitor.go
+++ b/janitor.go
@@ -27,13 +27,17 @@ type janitor struct {
 
 	// average interval between checks.
 	avgInterval time.Duration
+
+	// number of tasks to be deleted when janitor runs to delete the expired completed tasks.
+	batchSize int
 }
 
 type janitorParams struct {
-	logger   *log.Logger
-	broker   base.Broker
-	queues   []string
-	interval time.Duration
+	logger    *log.Logger
+	broker    base.Broker
+	queues    []string
+	interval  time.Duration
+	batchSize int
 }
 
 func newJanitor(params janitorParams) *janitor {
@@ -43,6 +47,7 @@ func newJanitor(params janitorParams) *janitor {
 		done:        make(chan struct{}),
 		queues:      params.queues,
 		avgInterval: params.interval,
+		batchSize:   params.batchSize,
 	}
 }
 
@@ -73,7 +78,7 @@ func (j *janitor) start(wg *sync.WaitGroup) {
 
 func (j *janitor) exec() {
 	for _, qname := range j.queues {
-		if err := j.broker.DeleteExpiredCompletedTasks(qname); err != nil {
+		if err := j.broker.DeleteExpiredCompletedTasks(qname, j.batchSize); err != nil {
 			j.logger.Errorf("Failed to delete expired completed tasks from queue %q: %v",
 				qname, err)
 		}
diff --git a/processor.go b/processor.go
index 0ba9890..4c6471a 100644
--- a/processor.go
+++ b/processor.go
@@ -415,21 +415,19 @@ func (p *processor) queues() []string {
 func (p *processor) perform(ctx context.Context, task *Task) (err error) {
 	defer func() {
 		if x := recover(); x != nil {
-			errMsg := string(debug.Stack())
-
-			p.logger.Errorf("recovering from panic. See the stack trace below for details:\n%s", errMsg)
+			p.logger.Errorf("recovering from panic. See the stack trace below for details:\n%s", string(debug.Stack()))
 			_, file, line, ok := runtime.Caller(1) // skip the first frame (panic itself)
 			if ok && strings.Contains(file, "runtime/") {
 				// The panic came from the runtime, most likely due to incorrect
 				// map/slice usage. The parent frame should have the real trigger.
 				_, file, line, ok = runtime.Caller(2)
 			}
-
+			var errMsg string
 			// Include the file and line number info in the error, if runtime.Caller returned ok.
 			if ok {
-				err = fmt.Errorf("panic [%s:%d]: %v", file, line, x)
+				errMsg = fmt.Sprintf("panic [%s:%d]: %v", file, line, x)
 			} else {
-				err = fmt.Errorf("panic: %v", x)
+				errMsg = fmt.Sprintf("panic: %v", x)
 			}
 			err = &errors.PanicError{
 				ErrMsg: errMsg,
diff --git a/scheduler.go b/scheduler.go
index 2655a94..5f0b05d 100644
--- a/scheduler.go
+++ b/scheduler.go
@@ -10,11 +10,11 @@ import (
 	"sync"
 	"time"
 
-	"github.com/redis/go-redis/v9"
 	"github.com/google/uuid"
 	"github.com/hibiken/asynq/internal/base"
 	"github.com/hibiken/asynq/internal/log"
 	"github.com/hibiken/asynq/internal/rdb"
+	"github.com/redis/go-redis/v9"
 	"github.com/robfig/cron/v3"
 )
 
@@ -43,15 +43,27 @@ type Scheduler struct {
 	// to avoid using cron.EntryID as the public API of
 	// the Scheduler.
 	idmap map[string]cron.EntryID
+	// When a Scheduler has been created with an existing Redis connection, we do
+	// not want to close it.
+	sharedConnection bool
 }
 
 // NewScheduler returns a new Scheduler instance given the redis connection option.
 // The parameter opts is optional, defaults will be used if opts is set to nil
 func NewScheduler(r RedisConnOpt, opts *SchedulerOpts) *Scheduler {
-	c, ok := r.MakeRedisClient().(redis.UniversalClient)
+	redisClient, ok := r.MakeRedisClient().(redis.UniversalClient)
 	if !ok {
 		panic(fmt.Sprintf("asynq: unsupported RedisConnOpt type %T", r))
 	}
+	scheduler := NewSchedulerFromRedisClient(redisClient, opts)
+	scheduler.sharedConnection = false
+	return scheduler
+}
+
+// NewSchedulerFromRedisClient returns a new instance of Scheduler given a redis.UniversalClient
+// The parameter opts is optional, defaults will be used if opts is set to nil.
+// Warning: The underlying redis connection pool will not be closed by Asynq, you are responsible for closing it.
+func NewSchedulerFromRedisClient(c redis.UniversalClient, opts *SchedulerOpts) *Scheduler {
 	if opts == nil {
 		opts = &SchedulerOpts{}
 	}
@@ -72,7 +84,7 @@ func NewScheduler(r RedisConnOpt, opts *SchedulerOpts) *Scheduler {
 		id:              generateSchedulerID(),
 		state:           &serverState{value: srvStateNew},
 		logger:          logger,
-		client:          NewClient(r),
+		client:          NewClientFromRedisClient(c),
 		rdb:             rdb.NewRDB(c),
 		cron:            cron.New(cron.WithLocation(loc)),
 		location:        loc,
@@ -261,8 +273,12 @@ func (s *Scheduler) Shutdown() {
 	s.wg.Wait()
 
 	s.clearHistory()
-	s.client.Close()
-	s.rdb.Close()
+	if err := s.client.Close(); err != nil {
+		s.logger.Errorf("Failed to close redis client connection: %v", err)
+	}
+	if !s.sharedConnection {
+		s.rdb.Close()
+	}
 	s.logger.Info("Scheduler stopped")
 }
 
@@ -320,3 +336,14 @@ func (s *Scheduler) clearHistory() {
 		}
 	}
 }
+
+// Ping performs a ping against the redis connection.
+func (s *Scheduler) Ping() error {
+	s.state.mu.Lock()
+	defer s.state.mu.Unlock()
+	if s.state.value == srvStateClosed {
+		return nil
+	}
+
+	return s.rdb.Ping()
+}
diff --git a/server.go b/server.go
index cc6f418..e405399 100644
--- a/server.go
+++ b/server.go
@@ -37,6 +37,9 @@ type Server struct {
 	logger *log.Logger
 
 	broker base.Broker
+	// When a Server has been created with an existing Redis connection, we do
+	// not want to close it.
+	sharedConnection bool
 
 	state *serverState
 
@@ -103,7 +106,7 @@ type Config struct {
 	// If BaseContext is nil, the default is context.Background().
 	// If this is defined, then it MUST return a non-nil context
 	BaseContext func() context.Context
-	
+
 	// TaskCheckInterval specifies the interval between checks for new tasks to process when all queues are empty.
 	//
 	// If unset, zero or a negative value, the interval is set to 1 second.
@@ -239,6 +242,17 @@ type Config struct {
 	//
 	// If unset or nil, the group aggregation feature will be disabled on the server.
 	GroupAggregator GroupAggregator
+
+	// JanitorInterval specifies the average interval of janitor checks for expired completed tasks.
+	//
+	// If unset or zero, default interval of 8 seconds is used.
+	JanitorInterval time.Duration
+
+	// JanitorBatchSize specifies the number of expired completed tasks to be deleted in one run.
+	//
+	// If unset or zero, default batch size of 100 is used.
+	// Make sure to not put a big number as the batch size to prevent a long-running script.
+	JanitorBatchSize int
 }
 
 // GroupAggregator aggregates a group of tasks into one before the tasks are passed to the Handler.
@@ -408,15 +422,28 @@ const (
 	defaultDelayedTaskCheckInterval = 5 * time.Second
 
 	defaultGroupGracePeriod = 1 * time.Minute
+
+	defaultJanitorInterval = 8 * time.Second
+
+	defaultJanitorBatchSize = 100
 )
 
 // NewServer returns a new Server given a redis connection option
 // and server configuration.
 func NewServer(r RedisConnOpt, cfg Config) *Server {
-	c, ok := r.MakeRedisClient().(redis.UniversalClient)
+	redisClient, ok := r.MakeRedisClient().(redis.UniversalClient)
 	if !ok {
 		panic(fmt.Sprintf("asynq: unsupported RedisConnOpt type %T", r))
 	}
+	server := NewServerFromRedisClient(redisClient, cfg)
+	server.sharedConnection = false
+	return server
+}
+
+// NewServerFromRedisClient returns a new instance of Server given a redis.UniversalClient
+// and server configuration
+// Warning: The underlying redis connection pool will not be closed by Asynq, you are responsible for closing it.
+func NewServerFromRedisClient(c redis.UniversalClient, cfg Config) *Server {
 	baseCtxFn := cfg.BaseContext
 	if baseCtxFn == nil {
 		baseCtxFn = context.Background
@@ -547,11 +574,26 @@ func NewServer(r RedisConnOpt, cfg Config) *Server {
 		interval:        healthcheckInterval,
 		healthcheckFunc: cfg.HealthCheckFunc,
 	})
+
+	janitorInterval := cfg.JanitorInterval
+	if janitorInterval == 0 {
+		janitorInterval = defaultJanitorInterval
+	}
+
+	janitorBatchSize := cfg.JanitorBatchSize
+	if janitorBatchSize == 0 {
+		janitorBatchSize = defaultJanitorBatchSize
+	}
+	if janitorBatchSize > defaultJanitorBatchSize {
+		logger.Warnf("Janitor batch size of %d is greater than the recommended batch size of %d. "+
+			"This might cause a long-running script", janitorBatchSize, defaultJanitorBatchSize)
+	}
 	janitor := newJanitor(janitorParams{
-		logger:   logger,
-		broker:   rdb,
-		queues:   qnames,
-		interval: 8 * time.Second,
+		logger:    logger,
+		broker:    rdb,
+		queues:    qnames,
+		interval:  janitorInterval,
+		batchSize: janitorBatchSize,
 	})
 	aggregator := newAggregator(aggregatorParams{
 		logger:          logger,
@@ -563,18 +605,19 @@ func NewServer(r RedisConnOpt, cfg Config) *Server {
 		groupAggregator: cfg.GroupAggregator,
 	})
 	return &Server{
-		logger:        logger,
-		broker:        rdb,
-		state:         srvState,
-		forwarder:     forwarder,
-		processor:     processor,
-		syncer:        syncer,
-		heartbeater:   heartbeater,
-		subscriber:    subscriber,
-		recoverer:     recoverer,
-		healthchecker: healthchecker,
-		janitor:       janitor,
-		aggregator:    aggregator,
+		logger:           logger,
+		broker:           rdb,
+		sharedConnection: true,
+		state:            srvState,
+		forwarder:        forwarder,
+		processor:        processor,
+		syncer:           syncer,
+		heartbeater:      heartbeater,
+		subscriber:       subscriber,
+		recoverer:        recoverer,
+		healthchecker:    healthchecker,
+		janitor:          janitor,
+		aggregator:       aggregator,
 	}
 }
 
@@ -702,7 +745,9 @@ func (srv *Server) Shutdown() {
 	srv.heartbeater.shutdown()
 	srv.wg.Wait()
 
-	srv.broker.Close()
+	if !srv.sharedConnection {
+		srv.broker.Close()
+	}
 	srv.logger.Info("Exiting")
 }
 
@@ -714,7 +759,7 @@ func (srv *Server) Shutdown() {
 func (srv *Server) Stop() {
 	srv.state.mu.Lock()
 	if srv.state.value != srvStateActive {
-		// Invalid calll to Stop, server can only go from Active state to Stopped state.
+		// Invalid call to Stop, server can only go from Active state to Stopped state.
 		srv.state.mu.Unlock()
 		return
 	}
@@ -725,3 +770,16 @@ func (srv *Server) Stop() {
 	srv.processor.stop()
 	srv.logger.Info("Processor stopped")
 }
+
+// Ping performs a ping against the redis connection.
+//
+// This is an alternative to the HealthCheckFunc available in the Config object.
+func (srv *Server) Ping() error {
+	srv.state.mu.Lock()
+	defer srv.state.mu.Unlock()
+	if srv.state.value == srvStateClosed {
+		return nil
+	}
+
+	return srv.broker.Ping()
+}


Test patch:
diff --git a/client_test.go b/client_test.go
index da24d13..80a5517 100644
--- a/client_test.go
+++ b/client_test.go
@@ -14,6 +14,7 @@ import (
 	"github.com/google/go-cmp/cmp/cmpopts"
 	"github.com/hibiken/asynq/internal/base"
 	h "github.com/hibiken/asynq/internal/testutil"
+	"github.com/redis/go-redis/v9"
 )
 
 func TestClientEnqueueWithProcessAtOption(t *testing.T) {
@@ -143,11 +144,7 @@ func TestClientEnqueueWithProcessAtOption(t *testing.T) {
 	}
 }
 
-func TestClientEnqueue(t *testing.T) {
-	r := setup(t)
-	client := NewClient(getRedisConnOpt(t))
-	defer client.Close()
-
+func testClientEnqueue(t *testing.T, client *Client, r redis.UniversalClient) {
 	task := NewTask("send_email", h.JSON(map[string]interface{}{"to": "customer@gmail.com", "from": "merchant@example.com"}))
 	now := time.Now()
 
@@ -478,6 +475,24 @@ func TestClientEnqueue(t *testing.T) {
 	}
 }
 
+func TestClientEnqueue(t *testing.T) {
+	r := setup(t)
+	client := NewClient(getRedisConnOpt(t))
+	defer client.Close()
+	testClientEnqueue(t, client, r)
+}
+
+func TestClientFromRedisClientEnqueue(t *testing.T) {
+	r := setup(t)
+	redisClient := getRedisConnOpt(t).MakeRedisClient().(redis.UniversalClient)
+	client := NewClientFromRedisClient(redisClient)
+	testClientEnqueue(t, client, r)
+	err := client.Close()
+	if err == nil {
+		t.Error("client.Close() should have failed because of a shared client but it didn't")
+	}
+}
+
 func TestClientEnqueueWithGroupOption(t *testing.T) {
 	r := setup(t)
 	client := NewClient(getRedisConnOpt(t))
@@ -541,11 +556,11 @@ func TestClientEnqueueWithGroupOption(t *testing.T) {
 			},
 		},
 		{
-			desc: "With Group and ProcessIn options",
+			desc: "With Group and ProcessAt options",
 			task: task,
 			opts: []Option{
 				Group("mygroup"),
-				ProcessIn(30 * time.Minute),
+				ProcessAt(now.Add(30 * time.Minute)),
 			},
 			wantInfo: &TaskInfo{
 				Queue:         "default",
diff --git a/inspector_test.go b/inspector_test.go
index 0d47d06..1d30a03 100644
--- a/inspector_test.go
+++ b/inspector_test.go
@@ -22,11 +22,7 @@ import (
 	"github.com/redis/go-redis/v9"
 )
 
-func TestInspectorQueues(t *testing.T) {
-	r := setup(t)
-	defer r.Close()
-	inspector := NewInspector(getRedisConnOpt(t))
-
+func testInspectorQueues(t *testing.T, inspector *Inspector, r redis.UniversalClient) {
 	tests := []struct {
 		queues []string
 	}{
@@ -52,7 +48,21 @@ func TestInspectorQueues(t *testing.T) {
 			t.Errorf("Queues() = %v, want %v; (-want, +got)\n%s", got, tc.queues, diff)
 		}
 	}
+}
+
+func TestInspectorQueues(t *testing.T) {
+	r := setup(t)
+	defer r.Close()
+	inspector := NewInspector(getRedisConnOpt(t))
+	testInspectorQueues(t, inspector, r)
+}
 
+func TestInspectorFromRedisClientQueues(t *testing.T) {
+	r := setup(t)
+	defer r.Close()
+	redisClient := getRedisConnOpt(t).MakeRedisClient().(redis.UniversalClient)
+	inspector := NewInspectorFromRedisClient(redisClient)
+	testInspectorQueues(t, inspector, r)
 }
 
 func TestInspectorDeleteQueue(t *testing.T) {
diff --git a/internal/rdb/rdb_test.go b/internal/rdb/rdb_test.go
index 3bd9eda..1b7d138 100644
--- a/internal/rdb/rdb_test.go
+++ b/internal/rdb/rdb_test.go
@@ -2002,7 +2002,6 @@ func TestArchive(t *testing.T) {
 	}
 	errMsg := "SMTP server not responding"
 
-	// TODO(hibiken): add test cases for trimming
 	tests := []struct {
 		active       map[string][]*base.TaskMessage
 		lease        map[string][]base.Z
@@ -2171,6 +2170,163 @@ func TestArchive(t *testing.T) {
 	}
 }
 
+func TestArchiveTrim(t *testing.T) {
+	r := setup(t)
+	defer r.Close()
+	now := time.Now()
+	r.SetClock(timeutil.NewSimulatedClock(now))
+
+	t1 := &base.TaskMessage{
+		ID:      uuid.NewString(),
+		Type:    "send_email",
+		Payload: nil,
+		Queue:   "default",
+		Retry:   25,
+		Retried: 25,
+		Timeout: 1800,
+	}
+	t2 := &base.TaskMessage{
+		ID:      uuid.NewString(),
+		Type:    "reindex",
+		Payload: nil,
+		Queue:   "default",
+		Retry:   25,
+		Retried: 0,
+		Timeout: 3000,
+	}
+	errMsg := "SMTP server not responding"
+
+	maxArchiveSet := make([]base.Z, 0)
+	for i := 0; i < maxArchiveSize-1; i++ {
+		maxArchiveSet = append(maxArchiveSet, base.Z{Message: &base.TaskMessage{
+			ID:      uuid.NewString(),
+			Type:    "generate_csv",
+			Payload: nil,
+			Queue:   "default",
+			Retry:   25,
+			Retried: 0,
+			Timeout: 60,
+		}, Score: now.Add(-time.Hour + -time.Second*time.Duration(i)).Unix()})
+	}
+
+	wantMaxArchiveSet := make([]base.Z, 0)
+	// newly archived task should be at the front
+	wantMaxArchiveSet = append(wantMaxArchiveSet, base.Z{Message: h.TaskMessageWithError(*t1, errMsg, now), Score: now.Unix()})
+	// oldest task should be dropped from the set
+	wantMaxArchiveSet = append(wantMaxArchiveSet, maxArchiveSet[:len(maxArchiveSet)-1]...)
+
+	tests := []struct {
+		toArchive    map[string][]*base.TaskMessage
+		lease        map[string][]base.Z
+		archived     map[string][]base.Z
+		wantArchived map[string][]base.Z
+	}{
+		{ // simple, 1 to be archived, 1 already archived, both are in the archive set
+			toArchive: map[string][]*base.TaskMessage{
+				"default": {t1},
+			},
+			lease: map[string][]base.Z{
+				"default": {
+					{Message: t1, Score: now.Add(10 * time.Second).Unix()},
+				},
+			},
+			archived: map[string][]base.Z{
+				"default": {
+					{Message: t2, Score: now.Add(-time.Hour).Unix()},
+				},
+			},
+			wantArchived: map[string][]base.Z{
+				"default": {
+					{Message: h.TaskMessageWithError(*t1, errMsg, now), Score: now.Unix()},
+					{Message: t2, Score: now.Add(-time.Hour).Unix()},
+				},
+			},
+		},
+		{ // 1 to be archived, 1 already archived but past expiry, only the newly archived task should be left
+			toArchive: map[string][]*base.TaskMessage{
+				"default": {t1},
+			},
+			lease: map[string][]base.Z{
+				"default": {
+					{Message: t1, Score: now.Add(10 * time.Second).Unix()},
+				},
+			},
+			archived: map[string][]base.Z{
+				"default": {
+					{Message: t2, Score: now.Add(-time.Hour * 24 * (archivedExpirationInDays + 1)).Unix()},
+				},
+			},
+			wantArchived: map[string][]base.Z{
+				"default": {
+					{Message: h.TaskMessageWithError(*t1, errMsg, now), Score: now.Unix()},
+				},
+			},
+		},
+		{ // 1 to be archived, maxArchiveSize in archive set, archive set should be trimmed back to maxArchiveSize and newly archived task should be in the set
+			toArchive: map[string][]*base.TaskMessage{
+				"default": {t1},
+			},
+			lease: map[string][]base.Z{
+				"default": {
+					{Message: t1, Score: now.Add(10 * time.Second).Unix()},
+				},
+			},
+			archived: map[string][]base.Z{
+				"default": maxArchiveSet,
+			},
+			wantArchived: map[string][]base.Z{
+				"default": wantMaxArchiveSet,
+			},
+		},
+	}
+
+	for _, tc := range tests {
+		h.FlushDB(t, r.client) // clean up db before each test case
+		h.SeedAllActiveQueues(t, r.client, tc.toArchive)
+		h.SeedAllLease(t, r.client, tc.lease)
+		h.SeedAllArchivedQueues(t, r.client, tc.archived)
+
+		for _, tasks := range tc.toArchive {
+			for _, target := range tasks {
+				err := r.Archive(context.Background(), target, errMsg)
+				if err != nil {
+					t.Errorf("(*RDB).Archive(%v, %v) = %v, want nil", target, errMsg, err)
+					continue
+				}
+			}
+		}
+
+		for queue, want := range tc.wantArchived {
+			gotArchived := h.GetArchivedEntries(t, r.client, queue)
+
+			if diff := cmp.Diff(want, gotArchived, h.SortZSetEntryOpt, zScoreCmpOpt, timeCmpOpt); diff != "" {
+				t.Errorf("mismatch found in %q after calling (*RDB).Archive: (-want, +got):\n%s", base.ArchivedKey(queue), diff)
+			}
+
+			// check that only keys present in the archived set are in rdb
+			vals := r.client.Keys(context.Background(), base.TaskKeyPrefix(queue)+"*").Val()
+			if len(vals) != len(gotArchived) {
+				t.Errorf("len of keys = %v, want %v", len(vals), len(gotArchived))
+				return
+			}
+
+			for _, val := range vals {
+				found := false
+				for _, entry := range gotArchived {
+					if strings.Contains(val, entry.Message.ID) {
+						found = true
+						break
+					}
+				}
+
+				if !found {
+					t.Errorf("key %v not found in archived set (it was orphaned by the archive trim)", val)
+				}
+			}
+		}
+	}
+}
+
 func TestForwardIfReadyWithGroup(t *testing.T) {
 	r := setup(t)
 	defer r.Close()
@@ -2542,8 +2698,8 @@ func TestDeleteExpiredCompletedTasks(t *testing.T) {
 		h.FlushDB(t, r.client)
 		h.SeedAllCompletedQueues(t, r.client, tc.completed)
 
-		if err := r.DeleteExpiredCompletedTasks(tc.qname); err != nil {
-			t.Errorf("DeleteExpiredCompletedTasks(%q) failed: %v", tc.qname, err)
+		if err := r.DeleteExpiredCompletedTasks(tc.qname, 100); err != nil {
+			t.Errorf("DeleteExpiredCompletedTasks(%q, 100) failed: %v", tc.qname, err)
 			continue
 		}
 
diff --git a/internal/testbroker/testbroker.go b/internal/testbroker/testbroker.go
index 368ae3e..ffab6fe 100644
--- a/internal/testbroker/testbroker.go
+++ b/internal/testbroker/testbroker.go
@@ -11,8 +11,8 @@ import (
 	"sync"
 	"time"
 
-	"github.com/redis/go-redis/v9"
 	"github.com/hibiken/asynq/internal/base"
+	"github.com/redis/go-redis/v9"
 )
 
 var errRedisDown = errors.New("testutil: redis is down")
@@ -145,13 +145,13 @@ func (tb *TestBroker) ForwardIfReady(qnames ...string) error {
 	return tb.real.ForwardIfReady(qnames...)
 }
 
-func (tb *TestBroker) DeleteExpiredCompletedTasks(qname string) error {
+func (tb *TestBroker) DeleteExpiredCompletedTasks(qname string, batchSize int) error {
 	tb.mu.Lock()
 	defer tb.mu.Unlock()
 	if tb.sleeping {
 		return errRedisDown
 	}
-	return tb.real.DeleteExpiredCompletedTasks(qname)
+	return tb.real.DeleteExpiredCompletedTasks(qname, batchSize)
 }
 
 func (tb *TestBroker) ListLeaseExpired(cutoff time.Time, qnames ...string) ([]*base.TaskMessage, error) {
diff --git a/janitor_test.go b/janitor_test.go
index 7f61278..a8bd02a 100644
--- a/janitor_test.go
+++ b/janitor_test.go
@@ -26,11 +26,13 @@ func TestJanitor(t *testing.T) {
 	defer r.Close()
 	rdbClient := rdb.NewRDB(r)
 	const interval = 1 * time.Second
+	const batchSize = 100
 	janitor := newJanitor(janitorParams{
-		logger:   testLogger,
-		broker:   rdbClient,
-		queues:   []string{"default", "custom"},
-		interval: interval,
+		logger:    testLogger,
+		broker:    rdbClient,
+		queues:    []string{"default", "custom"},
+		interval:  interval,
+		batchSize: batchSize,
 	})
 
 	now := time.Now()
diff --git a/scheduler_test.go b/scheduler_test.go
index fea048e..a35f390 100644
--- a/scheduler_test.go
+++ b/scheduler_test.go
@@ -5,6 +5,7 @@
 package asynq
 
 import (
+	"github.com/redis/go-redis/v9"
 	"sync"
 	"testing"
 	"time"
@@ -58,6 +59,7 @@ func TestSchedulerRegister(t *testing.T) {
 
 	r := setup(t)
 
+	// Tests for new redis connection.
 	for _, tc := range tests {
 		scheduler := NewScheduler(getRedisConnOpt(t), nil)
 		if _, err := scheduler.Register(tc.cronspec, tc.task, tc.opts...); err != nil {
@@ -75,6 +77,28 @@ func TestSchedulerRegister(t *testing.T) {
 			t.Errorf("mismatch found in queue %q: (-want,+got)\n%s", tc.queue, diff)
 		}
 	}
+
+	r = setup(t)
+
+	// Tests for existing redis connection.
+	for _, tc := range tests {
+		redisClient := getRedisConnOpt(t).MakeRedisClient().(redis.UniversalClient)
+		scheduler := NewSchedulerFromRedisClient(redisClient, nil)
+		if _, err := scheduler.Register(tc.cronspec, tc.task, tc.opts...); err != nil {
+			t.Fatal(err)
+		}
+
+		if err := scheduler.Start(); err != nil {
+			t.Fatal(err)
+		}
+		time.Sleep(tc.wait)
+		scheduler.Shutdown()
+
+		got := testutil.GetPendingMessages(t, r, tc.queue)
+		if diff := cmp.Diff(tc.want, got, testutil.IgnoreIDOpt); diff != "" {
+			t.Errorf("mismatch found in queue %q: (-want,+got)\n%s", tc.queue, diff)
+		}
+	}
 }
 
 func TestSchedulerWhenRedisDown(t *testing.T) {
diff --git a/server_test.go b/server_test.go
index 724b48a..da7aac5 100644
--- a/server_test.go
+++ b/server_test.go
@@ -14,22 +14,12 @@ import (
 	"github.com/hibiken/asynq/internal/rdb"
 	"github.com/hibiken/asynq/internal/testbroker"
 	"github.com/hibiken/asynq/internal/testutil"
+
+	"github.com/redis/go-redis/v9"
 	"go.uber.org/goleak"
 )
 
-func TestServer(t *testing.T) {
-	// https://github.com/go-redis/redis/issues/1029
-	ignoreOpt := goleak.IgnoreTopFunction("github.com/redis/go-redis/v9/internal/pool.(*ConnPool).reaper")
-	defer goleak.VerifyNone(t, ignoreOpt)
-
-	redisConnOpt := getRedisConnOpt(t)
-	c := NewClient(redisConnOpt)
-	defer c.Close()
-	srv := NewServer(redisConnOpt, Config{
-		Concurrency: 10,
-		LogLevel:    testLogLevel,
-	})
-
+func testServer(t *testing.T, c *Client, srv *Server) {
 	// no-op handler
 	h := func(ctx context.Context, task *Task) error {
 		return nil
@@ -53,6 +43,43 @@ func TestServer(t *testing.T) {
 	srv.Shutdown()
 }
 
+func TestServer(t *testing.T) {
+	// https://github.com/go-redis/redis/issues/1029
+	ignoreOpt := goleak.IgnoreTopFunction("github.com/redis/go-redis/v9/internal/pool.(*ConnPool).reaper")
+	defer goleak.VerifyNone(t, ignoreOpt)
+
+	redisConnOpt := getRedisConnOpt(t)
+	c := NewClient(redisConnOpt)
+	defer c.Close()
+	srv := NewServer(redisConnOpt, Config{
+		Concurrency: 10,
+		LogLevel:    testLogLevel,
+	})
+
+	testServer(t, c, srv)
+}
+
+func TestServerFromRedisClient(t *testing.T) {
+	// https://github.com/go-redis/redis/issues/1029
+	ignoreOpt := goleak.IgnoreTopFunction("github.com/redis/go-redis/v9/internal/pool.(*ConnPool).reaper")
+	defer goleak.VerifyNone(t, ignoreOpt)
+
+	redisConnOpt := getRedisConnOpt(t)
+	redisClient := redisConnOpt.MakeRedisClient().(redis.UniversalClient)
+	c := NewClientFromRedisClient(redisClient)
+	srv := NewServerFromRedisClient(redisClient, Config{
+		Concurrency: 10,
+		LogLevel:    testLogLevel,
+	})
+
+	testServer(t, c, srv)
+
+	err := c.Close()
+	if err == nil {
+		t.Error("client.Close() should have failed because of a shared client but it didn't")
+	}
+}
+
 func TestServerRun(t *testing.T) {
 	// https://github.com/go-redis/redis/issues/1029
 	ignoreOpt := goleak.IgnoreTopFunction("github.com/redis/go-redis/v9/internal/pool.(*ConnPool).reaper")