asynq/x/rate/semaphore_test.go

package rate

import (
	"context"
	"flag"
	"fmt"
	"strings"
	"testing"
	"time"

	"github.com/google/uuid"
	"github.com/hibiken/asynq"
	"github.com/hibiken/asynq/internal/base"
	asynqcontext "github.com/hibiken/asynq/internal/context"
	"github.com/redis/go-redis/v9"
)

var (
	redisAddr string
	redisDB   int

	useRedisCluster   bool
	redisClusterAddrs string // comma-separated list of host:port
)

func init() {
	flag.StringVar(&redisAddr, "redis_addr", "localhost:6379", "redis address to use in testing")
	flag.IntVar(&redisDB, "redis_db", 14, "redis db number to use in testing")
	flag.BoolVar(&useRedisCluster, "redis_cluster", false, "use redis cluster as a broker in testing")
	flag.StringVar(&redisClusterAddrs, "redis_cluster_addrs", "localhost:7000,localhost:7001,localhost:7002", "comma separated list of redis server addresses")
}

func TestNewSemaphore(t *testing.T) {
	tests := []struct {
		desc           string
		name           string
		maxConcurrency int
		wantPanic      string
		connOpt        asynq.RedisConnOpt
	}{
		{
			desc:      "Bad RedisConnOpt",
			wantPanic: "rate.NewSemaphore: unsupported RedisConnOpt type *rate.badConnOpt",
			connOpt:   &badConnOpt{},
		},
		{
			desc:      "Zero maxTokens should panic",
			wantPanic: "rate.NewSemaphore: maxTokens cannot be less than 1",
		},
		{
			desc:           "Empty scope should panic",
			maxConcurrency: 2,
			name:           "    ",
			wantPanic:      "rate.NewSemaphore: scope should not be empty",
		},
	}

	for _, tt := range tests {
		t.Run(tt.desc, func(t *testing.T) {
			if tt.wantPanic != "" {
				defer func() {
					if r := recover(); r.(string) != tt.wantPanic {
						t.Errorf("%s;\nNewSemaphore should panic with msg: %s, got %s", tt.desc, tt.wantPanic, r.(string))
					}
				}()
			}

			opt := tt.connOpt
			if tt.connOpt == nil {
				opt = getRedisConnOpt(t)
			}

			sema := NewSemaphore(opt, tt.name, tt.maxConcurrency)
			defer sema.Close()
		})
	}
}

func TestNewSemaphore_Acquire(t *testing.T) {
	tests := []struct {
		desc           string
		name           string
		maxConcurrency int
		taskIDs        []string
		ctxFunc        func(string) (context.Context, context.CancelFunc)
		want           []bool
	}{
		{
			desc:           "Should acquire token when current token count is less than maxTokens",
			name:           "task-1",
			maxConcurrency: 3,
			taskIDs:        []string{uuid.NewString(), uuid.NewString()},
			ctxFunc: func(id string) (context.Context, context.CancelFunc) {
				return asynqcontext.New(context.Background(), &base.TaskMessage{
					ID:    id,
					Queue: "task-1",
				}, time.Now().Add(time.Second))
			},
			want: []bool{true, true},
		},
		{
			desc:           "Should fail acquiring token when current token count is equal to maxTokens",
			name:           "task-2",
			maxConcurrency: 3,
			taskIDs:        []string{uuid.NewString(), uuid.NewString(), uuid.NewString(), uuid.NewString()},
			ctxFunc: func(id string) (context.Context, context.CancelFunc) {
				return asynqcontext.New(context.Background(), &base.TaskMessage{
					ID:    id,
					Queue: "task-2",
				}, time.Now().Add(time.Second))
			},
			want: []bool{true, true, true, false},
		},
	}

	for _, tt := range tests {
		t.Run(tt.desc, func(t *testing.T) {
			opt := getRedisConnOpt(t)
			rc := opt.MakeRedisClient().(redis.UniversalClient)
			defer rc.Close()

			if err := rc.Del(context.Background(), semaphoreKey(tt.name)).Err(); err != nil {
				t.Errorf("%s;\nredis.UniversalClient.Del() got error %v", tt.desc, err)
			}

			sema := NewSemaphore(opt, tt.name, tt.maxConcurrency)
			defer sema.Close()

			for i := 0; i < len(tt.taskIDs); i++ {
				ctx, cancel := tt.ctxFunc(tt.taskIDs[i])

				got, err := sema.Acquire(ctx)
				if err != nil {
					t.Errorf("%s;\nSemaphore.Acquire() got error %v", tt.desc, err)
				}

				if got != tt.want[i] {
					t.Errorf("%s;\nSemaphore.Acquire(ctx) returned %v, want %v", tt.desc, got, tt.want[i])
				}

				cancel()
			}
		})
	}
}

func TestNewSemaphore_Acquire_Error(t *testing.T) {
	tests := []struct {
		desc           string
		name           string
		maxConcurrency int
		taskIDs        []string
		ctxFunc        func(string) (context.Context, context.CancelFunc)
		errStr         string
	}{
		{
			desc:           "Should return error if context has no deadline",
			name:           "task-3",
			maxConcurrency: 1,
			taskIDs:        []string{uuid.NewString(), uuid.NewString()},
			ctxFunc: func(id string) (context.Context, context.CancelFunc) {
				return context.Background(), func() {}
			},
			errStr: "provided context must have a deadline",
		},
		{
			desc:           "Should return error when context is missing taskID",
			name:           "task-4",
			maxConcurrency: 1,
			taskIDs:        []string{uuid.NewString()},
			ctxFunc: func(_ string) (context.Context, context.CancelFunc) {
				return context.WithTimeout(context.Background(), time.Second)
			},
			errStr: "provided context is missing task ID value",
		},
	}

	for _, tt := range tests {
		t.Run(tt.desc, func(t *testing.T) {
			opt := getRedisConnOpt(t)
			rc := opt.MakeRedisClient().(redis.UniversalClient)
			defer rc.Close()

			if err := rc.Del(context.Background(), semaphoreKey(tt.name)).Err(); err != nil {
				t.Errorf("%s;\nredis.UniversalClient.Del() got error %v", tt.desc, err)
			}

			sema := NewSemaphore(opt, tt.name, tt.maxConcurrency)
			defer sema.Close()

			for i := 0; i < len(tt.taskIDs); i++ {
				ctx, cancel := tt.ctxFunc(tt.taskIDs[i])

				_, err := sema.Acquire(ctx)
				if err == nil || err.Error() != tt.errStr {
					t.Errorf("%s;\nSemaphore.Acquire() got error %v want error %v", tt.desc, err, tt.errStr)
				}

				cancel()
			}
		})
	}
}

func TestNewSemaphore_Acquire_StaleToken(t *testing.T) {
	opt := getRedisConnOpt(t)
	rc := opt.MakeRedisClient().(redis.UniversalClient)
	defer rc.Close()

	taskID := uuid.NewString()

	// adding a set member to mimic the case where token is acquired but the goroutine crashed,
	// in which case, the token will not be explicitly removed and should be present already
	rc.ZAdd(context.Background(), semaphoreKey("stale-token"), redis.Z{
		Score:  float64(time.Now().Add(-10 * time.Second).Unix()),
		Member: taskID,
	})

	sema := NewSemaphore(opt, "stale-token", 1)
	defer sema.Close()

	ctx, cancel := asynqcontext.New(context.Background(), &base.TaskMessage{
		ID:    taskID,
		Queue: "task-1",
	}, time.Now().Add(time.Second))
	defer cancel()

	got, err := sema.Acquire(ctx)
	if err != nil {
		t.Errorf("Acquire_StaleToken;\nSemaphore.Acquire() got error %v", err)
	}

	if !got {
		t.Error("Acquire_StaleToken;\nSemaphore.Acquire() got false want true")
	}
}

func TestNewSemaphore_Release(t *testing.T) {
	tests := []struct {
		desc      string
		name      string
		taskIDs   []string
		ctxFunc   func(string) (context.Context, context.CancelFunc)
		wantCount int64
	}{
		{
			desc:    "Should decrease token count",
			name:    "task-5",
			taskIDs: []string{uuid.NewString()},
			ctxFunc: func(id string) (context.Context, context.CancelFunc) {
				return asynqcontext.New(context.Background(), &base.TaskMessage{
					ID:    id,
					Queue: "task-3",
				}, time.Now().Add(time.Second))
			},
		},
		{
			desc:    "Should decrease token count by 2",
			name:    "task-6",
			taskIDs: []string{uuid.NewString(), uuid.NewString()},
			ctxFunc: func(id string) (context.Context, context.CancelFunc) {
				return asynqcontext.New(context.Background(), &base.TaskMessage{
					ID:    id,
					Queue: "task-4",
				}, time.Now().Add(time.Second))
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.desc, func(t *testing.T) {
			opt := getRedisConnOpt(t)
			rc := opt.MakeRedisClient().(redis.UniversalClient)
			defer rc.Close()

			if err := rc.Del(context.Background(), semaphoreKey(tt.name)).Err(); err != nil {
				t.Errorf("%s;\nredis.UniversalClient.Del() got error %v", tt.desc, err)
			}

			var members []redis.Z
			for i := 0; i < len(tt.taskIDs); i++ {
				members = append(members, redis.Z{
					Score:  float64(time.Now().Add(time.Duration(i) * time.Second).Unix()),
					Member: tt.taskIDs[i],
				})
			}
			if err := rc.ZAdd(context.Background(), semaphoreKey(tt.name), members...).Err(); err != nil {
				t.Errorf("%s;\nredis.UniversalClient.ZAdd() got error %v", tt.desc, err)
			}

			sema := NewSemaphore(opt, tt.name, 3)
			defer sema.Close()

			for i := 0; i < len(tt.taskIDs); i++ {
				ctx, cancel := tt.ctxFunc(tt.taskIDs[i])

				if err := sema.Release(ctx); err != nil {
					t.Errorf("%s;\nSemaphore.Release() got error %v", tt.desc, err)
				}

				cancel()
			}

			i, err := rc.ZCount(context.Background(), semaphoreKey(tt.name), "-inf", "+inf").Result()
			if err != nil {
				t.Errorf("%s;\nredis.UniversalClient.ZCount() got error %v", tt.desc, err)
			}

			if i != tt.wantCount {
				t.Errorf("%s;\nSemaphore.Release(ctx) didn't release token, got %v want 0", tt.desc, i)
			}
		})
	}
}

func TestNewSemaphore_Release_Error(t *testing.T) {
	testID := uuid.NewString()

	tests := []struct {
		desc    string
		name    string
		taskIDs []string
		ctxFunc func(string) (context.Context, context.CancelFunc)
		errStr  string
	}{
		{
			desc:    "Should return error when context is missing taskID",
			name:    "task-7",
			taskIDs: []string{uuid.NewString()},
			ctxFunc: func(_ string) (context.Context, context.CancelFunc) {
				return context.WithTimeout(context.Background(), time.Second)
			},
			errStr: "provided context is missing task ID value",
		},
		{
			desc:    "Should return error when context has taskID which never acquired token",
			name:    "task-8",
			taskIDs: []string{uuid.NewString()},
			ctxFunc: func(_ string) (context.Context, context.CancelFunc) {
				return asynqcontext.New(context.Background(), &base.TaskMessage{
					ID:    testID,
					Queue: "task-4",
				}, time.Now().Add(time.Second))
			},
			errStr: fmt.Sprintf("no token found for task %q", testID),
		},
	}

	for _, tt := range tests {
		t.Run(tt.desc, func(t *testing.T) {
			opt := getRedisConnOpt(t)
			rc := opt.MakeRedisClient().(redis.UniversalClient)
			defer rc.Close()

			if err := rc.Del(context.Background(), semaphoreKey(tt.name)).Err(); err != nil {
				t.Errorf("%s;\nredis.UniversalClient.Del() got error %v", tt.desc, err)
			}

			var members []redis.Z
			for i := 0; i < len(tt.taskIDs); i++ {
				members = append(members, redis.Z{
					Score:  float64(time.Now().Add(time.Duration(i) * time.Second).Unix()),
					Member: tt.taskIDs[i],
				})
			}
			if err := rc.ZAdd(context.Background(), semaphoreKey(tt.name), members...).Err(); err != nil {
				t.Errorf("%s;\nredis.UniversalClient.ZAdd() got error %v", tt.desc, err)
			}

			sema := NewSemaphore(opt, tt.name, 3)
			defer sema.Close()

			for i := 0; i < len(tt.taskIDs); i++ {
				ctx, cancel := tt.ctxFunc(tt.taskIDs[i])

				if err := sema.Release(ctx); err == nil || err.Error() != tt.errStr {
					t.Errorf("%s;\nSemaphore.Release() got error %v want error %v", tt.desc, err, tt.errStr)
				}

				cancel()
			}
		})
	}
}

func getRedisConnOpt(tb testing.TB) asynq.RedisConnOpt {
	tb.Helper()
	if useRedisCluster {
		addrs := strings.Split(redisClusterAddrs, ",")
		if len(addrs) == 0 {
			tb.Fatal("No redis cluster addresses provided. Please set addresses using --redis_cluster_addrs flag.")
		}
		return asynq.RedisClusterClientOpt{
			Addrs: addrs,
		}
	}
	return asynq.RedisClientOpt{
		Addr: redisAddr,
		DB:   redisDB,
	}
}

type badConnOpt struct {
}

func (b badConnOpt) MakeRedisClient() interface{} {
	return nil
}