ring_test.go

package queue

import (
	"context"
	"errors"
	"fmt"
	"log"
	"runtime"
	"sync"
	"testing"
	"time"

	"github.com/golang-queue/queue/core"
	"github.com/golang-queue/queue/job"
	"github.com/golang-queue/queue/mocks"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"go.uber.org/mock/gomock"
)

const (
	testMessageFoo  = "foo"
	testMessageTest = "test"
)

func TestMaxCapacity(t *testing.T) {
	w := NewRing(WithQueueSize(2))

	require.NoError(t, w.Queue(&mockMessage{}))
	require.NoError(t, w.Queue(&mockMessage{}))
	require.Error(t, w.Queue(&mockMessage{}))

	err := w.Queue(&mockMessage{})
	assert.Equal(t, ErrMaxCapacity, err)
}

func TestConcurrentMaxCapacity(t *testing.T) {
	for _, capacity := range []int{1, 2, 8} {
		t.Run(fmt.Sprintf("capacity_%d", capacity), func(t *testing.T) {
			const producers = 64
			w := NewRing(WithQueueSize(capacity))
			start := make(chan struct{})
			results := make(chan error, producers)
			var workers sync.WaitGroup

			for range producers {
				workers.Go(func() {
					<-start
					results <- w.Queue(&mockMessage{})
				})
			}
			close(start)
			workers.Wait()
			close(results)

			accepted := 0
			for err := range results {
				if err == nil {
					accepted++
				} else {
					require.ErrorIs(t, err, ErrMaxCapacity)
				}
			}
			require.Equal(t, capacity, accepted)
			require.Equal(t, capacity, w.count)

			for range capacity {
				_, err := w.Request()
				require.NoError(t, err)
			}
			_, err := w.Request()
			require.ErrorIs(t, err, ErrNoTaskInQueue)
			require.NoError(t, w.Queue(&mockMessage{}))
		})
	}
}

func TestConcurrentRequestsDuringShutdown(t *testing.T) {
	const (
		tasks     = 256
		consumers = 8
	)
	w := NewRing()
	for range tasks {
		require.NoError(t, w.Queue(&mockMessage{}))
	}
	// Set the shutdown state directly to isolate concurrent draining.
	w.stopFlag.Store(1)
	start := make(chan struct{})
	results := make(chan error, tasks+consumers)
	var workers sync.WaitGroup

	for range consumers {
		workers.Go(func() {
			<-start
			for {
				_, err := w.Request()
				results <- err
				if err != nil {
					return
				}
			}
		})
	}
	close(start)
	workers.Wait()
	close(results)

	drained := 0
	for err := range results {
		if err == nil {
			drained++
		} else {
			require.ErrorIs(t, err, ErrQueueHasBeenClosed)
		}
	}
	require.Equal(t, tasks, drained)
	task, err := w.Request()
	require.Nil(t, task)
	require.ErrorIs(t, err, ErrQueueHasBeenClosed)
	require.ErrorIs(t, w.Queue(&mockMessage{}), ErrQueueShutdown)
}

func TestCustomFuncAndWait(t *testing.T) {
	m := mockMessage{
		message: testMessageFoo,
	}
	w := NewRing(
		WithFn(func(ctx context.Context, m core.TaskMessage) error {
			time.Sleep(500 * time.Millisecond)
			return nil
		}),
	)
	q, err := NewQueue(
		WithWorker(w),
		WithWorkerCount(2),
		WithLogger(NewLogger()),
	)
	require.NoError(t, err)
	require.NoError(t, q.Queue(m))
	require.NoError(t, q.Queue(m))
	require.NoError(t, q.Queue(m))
	require.NoError(t, q.Queue(m))
	q.Start()
	time.Sleep(100 * time.Millisecond)
	assert.Equal(t, 2, int(q.metric.BusyWorkers()))
	time.Sleep(600 * time.Millisecond)
	q.Shutdown()
	q.Wait()
	// you will see the execute time > 1000ms
}

func TestEnqueueJobAfterShutdown(t *testing.T) {
	m := mockMessage{
		message: testMessageFoo,
	}
	w := NewRing()
	q, err := NewQueue(
		WithWorker(w),
		WithWorkerCount(2),
	)
	require.NoError(t, err)
	q.Start()
	time.Sleep(50 * time.Millisecond)
	q.Shutdown()
	// can't queue task after shutdown
	err = q.Queue(m)
	require.Error(t, err)
	assert.Equal(t, ErrQueueShutdown, err)
	q.Wait()
}

func TestJobReachTimeout(t *testing.T) {
	m := mockMessage{
		message: testMessageFoo,
	}
	w := NewRing(
		WithFn(func(ctx context.Context, m core.TaskMessage) error {
			for {
				select {
				case <-ctx.Done():
					log.Println("get data:", string(m.Payload()))
					if errors.Is(ctx.Err(), context.Canceled) {
						log.Println("queue has been shutdown and cancel the job")
					} else if errors.Is(ctx.Err(), context.DeadlineExceeded) {
						log.Println("job deadline exceeded")
					}
					return nil
				default:
				}
				time.Sleep(50 * time.Millisecond)
			}
		}),
	)
	q, err := NewQueue(
		WithWorker(w),
		WithWorkerCount(2),
	)
	require.NoError(t, err)
	require.NoError(t, q.Queue(m, job.AllowOption{Timeout: job.Time(30 * time.Millisecond)}))
	q.Start()
	time.Sleep(50 * time.Millisecond)
	q.Release()
}

func TestCancelJobAfterShutdown(t *testing.T) {
	m := mockMessage{
		message: testMessageFoo,
	}
	w := NewRing(
		WithLogger(NewEmptyLogger()),
		WithFn(func(ctx context.Context, m core.TaskMessage) error {
			for {
				select {
				case <-ctx.Done():
					log.Println("get data:", string(m.Payload()))
					if errors.Is(ctx.Err(), context.Canceled) {
						log.Println("queue has been shutdown and cancel the job")
					} else if errors.Is(ctx.Err(), context.DeadlineExceeded) {
						log.Println("job deadline exceeded")
					}
					return nil
				default:
				}
				time.Sleep(50 * time.Millisecond)
			}
		}),
	)
	q, err := NewQueue(
		WithWorker(w),
		WithWorkerCount(2),
	)
	require.NoError(t, err)
	require.NoError(t, q.Queue(m, job.AllowOption{Timeout: job.Time(100 * time.Millisecond)}))
	require.NoError(t, q.Queue(m, job.AllowOption{Timeout: job.Time(100 * time.Millisecond)}))
	q.Start()
	time.Sleep(10 * time.Millisecond)
	assert.Equal(t, int64(2), q.BusyWorkers())
	q.Release()
}

func TestGoroutineLeak(t *testing.T) {
	w := NewRing(
		WithLogger(NewLogger()),
		WithFn(func(ctx context.Context, m core.TaskMessage) error {
			for {
				select {
				case <-ctx.Done():
					if errors.Is(ctx.Err(), context.Canceled) {
						log.Println(
							"queue has been shutdown and cancel the job: " + string(m.Payload()),
						)
					} else if errors.Is(ctx.Err(), context.DeadlineExceeded) {
						log.Println("job deadline exceeded: " + string(m.Payload()))
					}
					return nil
				default:
					log.Println("get data:", string(m.Payload()))
					time.Sleep(50 * time.Millisecond)
					return nil
				}
			}
		}),
	)
	q, err := NewQueue(
		WithLogger(NewLogger()),
		WithWorker(w),
		WithWorkerCount(10),
	)
	require.NoError(t, err)
	for i := range 400 {
		m := mockMessage{
			message: fmt.Sprintf("new message: %d", i+1),
		}

		require.NoError(t, q.Queue(m))
	}

	q.Start()
	time.Sleep(1 * time.Second)
	q.Release()
	//nolint:forbidigo // intentional diagnostic output for this goroutine-leak check
	fmt.Println("number of goroutines:", runtime.NumGoroutine())
}

func TestGoroutinePanic(t *testing.T) {
	m := mockMessage{
		message: testMessageFoo,
	}
	w := NewRing(
		WithFn(func(ctx context.Context, m core.TaskMessage) error {
			panic("missing something")
		}),
	)
	q, err := NewQueue(
		WithWorker(w),
		WithWorkerCount(2),
	)
	require.NoError(t, err)
	require.NoError(t, q.Queue(m))
	q.Start()
	time.Sleep(10 * time.Millisecond)
	q.Release()
}

func TestIncreaseWorkerCount(t *testing.T) {
	w := NewRing(
		WithLogger(NewEmptyLogger()),
		WithFn(func(ctx context.Context, m core.TaskMessage) error {
			time.Sleep(500 * time.Millisecond)
			return nil
		}),
	)
	q, err := NewQueue(
		WithLogger(NewLogger()),
		WithWorker(w),
		WithWorkerCount(5),
	)
	require.NoError(t, err)

	for i := 1; i <= 10; i++ {
		m := mockMessage{
			message: fmt.Sprintf("new message: %d", i),
		}
		require.NoError(t, q.Queue(m))
	}

	q.Start()
	time.Sleep(100 * time.Millisecond)
	assert.Equal(t, int64(5), q.BusyWorkers())
	q.UpdateWorkerCount(10)
	time.Sleep(100 * time.Millisecond)
	assert.Equal(t, int64(10), q.BusyWorkers())
	q.Release()
}

func TestDecreaseWorkerCount(t *testing.T) {
	w := NewRing(
		WithFn(func(ctx context.Context, m core.TaskMessage) error {
			time.Sleep(100 * time.Millisecond)
			return nil
		}),
	)
	q, err := NewQueue(
		WithLogger(NewLogger()),
		WithWorker(w),
		WithWorkerCount(5),
	)
	require.NoError(t, err)

	for i := 1; i <= 10; i++ {
		m := mockMessage{
			message: fmt.Sprintf("test message: %d", i),
		}
		require.NoError(t, q.Queue(m))
	}

	q.Start()
	time.Sleep(20 * time.Millisecond)
	assert.Equal(t, int64(5), q.BusyWorkers())
	q.UpdateWorkerCount(3)
	time.Sleep(100 * time.Millisecond)
	assert.Equal(t, int64(3), q.BusyWorkers())
	time.Sleep(100 * time.Millisecond)
	assert.Equal(t, int64(2), q.BusyWorkers())
	q.Release()
}

func TestHandleAllJobBeforeShutdownRing(t *testing.T) {
	controller := gomock.NewController(t)
	defer controller.Finish()

	m := mocks.NewMockTaskMessage(controller)

	w := NewRing(
		WithFn(func(ctx context.Context, m core.TaskMessage) error {
			time.Sleep(10 * time.Millisecond)
			return nil
		}),
	)

	done := make(chan struct{})
	require.NoError(t, w.Queue(m))
	require.NoError(t, w.Queue(m))
	go func() {
		assert.NoError(t, w.Shutdown())
		done <- struct{}{}
	}()
	time.Sleep(50 * time.Millisecond)
	task, err := w.Request()
	assert.NotNil(t, task)
	require.NoError(t, err)
	task, err = w.Request()
	assert.NotNil(t, task)
	require.NoError(t, err)
	task, err = w.Request()
	assert.Nil(t, task)
	require.ErrorIs(t, err, ErrQueueHasBeenClosed)
	<-done
}

func TestHandleAllJobBeforeShutdownRingInQueue(t *testing.T) {
	controller := gomock.NewController(t)
	defer controller.Finish()

	m := mocks.NewMockTaskMessage(controller)
	m.EXPECT().Bytes().Return([]byte("test")).AnyTimes()
	m.EXPECT().Payload().Return([]byte("test")).AnyTimes()

	messages := make(chan string, 10)

	w := NewRing(
		WithFn(func(ctx context.Context, m core.TaskMessage) error {
			time.Sleep(10 * time.Millisecond)
			messages <- string(m.Payload())
			return nil
		}),
	)

	q, err := NewQueue(
		WithLogger(NewLogger()),
		WithWorker(w),
		WithWorkerCount(1),
	)
	require.NoError(t, err)

	require.NoError(t, q.Queue(m))
	require.NoError(t, q.Queue(m))
	assert.Empty(t, messages)
	q.Start()
	q.Release()
	assert.Len(t, messages, 2)
}

func TestRetryCountWithNewMessage(t *testing.T) {
	controller := gomock.NewController(t)
	defer controller.Finish()

	m := mocks.NewMockQueuedMessage(controller)
	m.EXPECT().Bytes().Return([]byte("test")).AnyTimes()

	messages := make(chan string, 10)
	keep := make(chan struct{})
	count := 1

	w := NewRing(
		WithFn(func(ctx context.Context, m core.TaskMessage) error {
			if count%3 != 0 {
				count++
				return errors.New("count not correct")
			}
			close(keep)
			messages <- string(m.Payload())
			return nil
		}),
	)

	q, err := NewQueue(
		WithLogger(NewLogger()),
		WithWorker(w),
		WithWorkerCount(1),
	)
	require.NoError(t, err)

	require.NoError(t, q.Queue(
		m,
		job.AllowOption{
			RetryCount: job.Int64(3),
			RetryDelay: job.Time(50 * time.Millisecond),
		},
	))
	assert.Empty(t, messages)
	q.Start()
	// wait retry twice.
	<-keep
	q.Release()
	assert.Len(t, messages, 1)
}

func TestRetryCountWithNewTask(t *testing.T) {
	messages := make(chan string, 10)
	count := 1

	w := NewRing()

	q, err := NewQueue(
		WithLogger(NewLogger()),
		WithWorker(w),
		WithWorkerCount(1),
	)
	require.NoError(t, err)

	keep := make(chan struct{})

	require.NoError(t, q.QueueTask(
		func(ctx context.Context) error {
			if count%3 != 0 {
				count++
				return errors.New("count not correct")
			}
			close(keep)
			messages <- "foobar"
			return nil
		},
		job.AllowOption{
			RetryCount: job.Int64(3),
		},
	))
	assert.Empty(t, messages)
	q.Start()
	// wait retry twice.
	<-keep
	q.Release()
	assert.Len(t, messages, 1)
}

func TestCancelRetryCountWithNewTask(t *testing.T) {
	messages := make(chan string, 10)
	count := 1

	w := NewRing()

	q, err := NewQueue(
		WithLogger(NewLogger()),
		WithWorker(w),
		WithWorkerCount(1),
	)
	require.NoError(t, err)

	require.NoError(t, q.QueueTask(
		func(ctx context.Context) error {
			if count%3 != 0 {
				count++
				q.logger.Info("add count")
				return errors.New("count not correct")
			}
			messages <- "foobar"
			return nil
		},
		job.AllowOption{
			RetryCount: job.Int64(3),
			RetryDelay: job.Time(100 * time.Millisecond),
		},
	))
	assert.Empty(t, messages)
	q.Start()
	time.Sleep(50 * time.Millisecond)
	q.Release()
	assert.Empty(t, messages)
	assert.Equal(t, 2, count)
}

func TestCancelRetryCountWithNewMessage(t *testing.T) {
	controller := gomock.NewController(t)
	defer controller.Finish()

	m := mocks.NewMockQueuedMessage(controller)
	m.EXPECT().Bytes().Return([]byte("test")).AnyTimes()

	messages := make(chan string, 10)
	count := 1

	w := NewRing(
		WithFn(func(ctx context.Context, m core.TaskMessage) error {
			if count%3 != 0 {
				count++
				return errors.New("count not correct")
			}
			messages <- string(m.Payload())
			return nil
		}),
	)

	q, err := NewQueue(
		WithLogger(NewLogger()),
		WithWorker(w),
		WithWorkerCount(1),
	)
	require.NoError(t, err)

	require.NoError(t, q.Queue(
		m,
		job.AllowOption{
			RetryCount: job.Int64(3),
			RetryDelay: job.Time(100 * time.Millisecond),
		},
	))
	assert.Empty(t, messages)
	q.Start()
	time.Sleep(50 * time.Millisecond)
	q.Release()
	assert.Empty(t, messages)
	assert.Equal(t, 2, count)
}

func TestErrNoTaskInQueue(t *testing.T) {
	w := NewRing(
		WithFn(func(ctx context.Context, m core.TaskMessage) error {
			return nil
		}),
	)
	task, err := w.Request()
	assert.Nil(t, task)
	require.Error(t, err)
	assert.Equal(t, ErrNoTaskInQueue, err)
}

func BenchmarkRingQueue(b *testing.B) {
	b.Run("queue and request operations", func(b *testing.B) {
		w := NewRing(WithQueueSize(1000))
		m := mockMessage{message: testMessageTest}

		b.ResetTimer()
		b.RunParallel(func(pb *testing.PB) {
			for pb.Next() {
				_ = w.Queue(&m)
				_, _ = w.Request()
			}
		})
	})

	b.Run("concurrent queue operations", func(b *testing.B) {
		w := NewRing(WithQueueSize(1000))
		m := mockMessage{message: testMessageTest}

		b.ResetTimer()
		b.RunParallel(func(pb *testing.PB) {
			for pb.Next() {
				_ = w.Queue(&m)
			}
		})
	})

	b.Run("resize operations", func(b *testing.B) {
		w := NewRing()
		m := mockMessage{message: testMessageTest}

		b.ResetTimer()
		for i := 0; i < b.N; i++ {
			for range 100 {
				_ = w.Queue(&m)
			}
			for range 100 {
				_, _ = w.Request()
			}
		}
	})
}