test: stop active tasks before server shutdown

internal/timeutil/sleep.go

@@ -0,0 +1,19 @@
+package timeutil
+
+import (
+	"context"
+	"time"
+)
+
+// Sleep will wait for the specified duration or return on context
+// expiration.
+func Sleep(ctx context.Context, d time.Duration) error {
+	t := time.NewTimer(d)
+	select {
+	case <-ctx.Done():
+		t.Stop()
+		return ctx.Err()
+	case <-t.C:
+		return nil
+	}
+}

internal/timeutil/sleep_test.go

@@ -0,0 +1,47 @@
+package timeutil
+
+import (
+	"context"
+	"github.com/stretchr/testify/require"
+	"sync"
+	"testing"
+	"time"
+)
+
+func TestSleep(t *testing.T) {
+	tests := []struct {
+		name    string
+		timeout time.Duration
+		sleep   time.Duration
+		wantErr error
+	}{
+		{
+			name:    "success",
+			timeout: 30 * time.Millisecond,
+			sleep:   10 * time.Millisecond,
+			wantErr: nil,
+		},
+		{
+			name:    "timeout",
+			timeout: 10 * time.Millisecond,
+			sleep:   30 * time.Millisecond,
+			wantErr: context.DeadlineExceeded,
+		},
+	}
+	for _, tc := range tests {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			wg := sync.WaitGroup{}
+			ctx, cancel := context.WithTimeout(context.Background(), tc.timeout)
+			wg.Add(1)
+			go func() {
+				defer wg.Done()
+				err := Sleep(ctx, tc.sleep)
+				require.ErrorIs(t, tc.wantErr, err)
+			}()
+			time.Sleep(20 * time.Millisecond)
+			cancel()
+			wg.Wait()
+		})
+	}
+}

server_test.go

@@ -7,6 +7,8 @@ package asynq
 import (
 	"context"
 	"fmt"
+	"github.com/hibiken/asynq/internal/timeutil"
+	"github.com/redis/go-redis/v9"
 	"syscall"
 	"testing"
 	"time"
@@ -82,6 +84,48 @@ func TestServerRun(t *testing.T) {
 	}
 }
 
+func TestServerShutdown(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)
+	r, ok := redisConnOpt.MakeRedisClient().(redis.UniversalClient)
+	if !ok {
+		t.Fatalf("asynq: unsupported RedisConnOpt type %T", r)
+	}
+	testutil.FlushDB(t, r)
+
+	c := NewClient(redisConnOpt)
+	defer c.Close()
+	_, err := c.Enqueue(NewTask("send_email", testutil.JSON(map[string]interface{}{"recipient_id": 123})))
+	if err != nil {
+		t.Fatalf("could not enqueue a task: %v", err)
+	}
+
+	srv := NewServer(redisConnOpt, Config{LogLevel: testLogLevel, ShutdownTimeout: 6 * time.Second})
+	done := make(chan struct{})
+	mux := NewServeMux()
+	mux.HandleFunc("send_email", func(ctx context.Context, task *Task) error {
+		err := timeutil.Sleep(ctx, 10*time.Second)
+		time.Sleep(1 * time.Second)
+		close(done)
+		return err
+	})
+	if err = srv.Start(mux); err != nil {
+		t.Fatal(err)
+	}
+	time.Sleep(1 * time.Second)
+	srv.Shutdown()
+
+	// confirm that active tasks have stopped when server is shut down.
+	select {
+	case <-done:
+	default:
+		t.Error("active tasks were not stopped when the server was shut down")
+	}
+}
+
 func TestServerErrServerClosed(t *testing.T) {
 	srv := NewServer(RedisClientOpt{Addr: ":6379"}, Config{LogLevel: testLogLevel})
 	handler := NewServeMux()