xinmi/internal/router/sync_task_retry_test.go

package router

import (
	"context"
	"errors"
	"testing"
	"time"

	"github.com/Tencent/WeKnora/internal/types"
	"github.com/hibiken/asynq"
)

func TestSyncTaskExecutorInjectsRetryMetadata(t *testing.T) {
	executor := NewSyncTaskExecutor()
	observed := make(chan [2]int, 1)
	executor.RegisterHandler("test:retry-metadata", func(ctx context.Context, _ *asynq.Task) error {
		retried, maxRetry, ok := types.TaskRetryMetadataFromContext(ctx)
		if !ok {
			observed <- [2]int{-1, -1}
			return nil
		}
		observed <- [2]int{retried, maxRetry}
		return nil
	})

	task := asynq.NewTask("test:retry-metadata", nil)
	if _, err := executor.Enqueue(task, asynq.MaxRetry(3)); err != nil {
		t.Fatalf("enqueue: %v", err)
	}

	select {
	case got := <-observed:
		if got != [2]int{0, 3} {
			t.Fatalf("retry metadata = %v, want [0 3]", got)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("timed out waiting for sync task")
	}
}

func TestSyncTaskExecutorFencesReservedTaskID(t *testing.T) {
	executor := NewSyncTaskExecutor()
	started, release := make(chan struct{}), make(chan struct{})
	executor.RegisterHandler("test:reserved-id", func(context.Context, *asynq.Task) error {
		close(started)
		<-release
		return nil
	})
	task := asynq.NewTask("test:reserved-id", []byte(`{"knowledge_id":"document"}`))
	opts := []asynq.Option{asynq.TaskID("resumable:document"), asynq.Retention(time.Hour), asynq.MaxRetry(0)}
	info, err := executor.Enqueue(task, opts...)
	if err != nil || info.ID != "resumable:document" {
		t.Fatalf("reserved task: %v %v", info, err)
	}
	select {
	case <-started:
	case <-time.After(2 * time.Second):
		t.Fatal("task did not start")
	}
	if _, err := executor.Enqueue(task, opts...); !errors.Is(err, asynq.ErrTaskIDConflict) {
		t.Fatalf("concurrent reserved task was accepted: %v", err)
	}
	if !executor.HasActiveTaskID(info.ID) {
		t.Fatal("active reserved task was not detected")
	}
	if !executor.HasActiveKnowledgeTask("document") {
		t.Fatal("active knowledge task was not detected")
	}
	close(release)
	deadline := time.Now().Add(2 * time.Second)
	for time.Now().Before(deadline) {
		executor.mu.RLock()
		expiry := executor.taskIDs[info.ID]
		executor.mu.RUnlock()
		if !expiry.IsZero() {
			if executor.HasActiveTaskID(info.ID) {
				t.Fatal("retained task reported active")
			}
			if executor.HasActiveKnowledgeTask("document") {
				t.Fatal("completed knowledge task reported active")
			}
			if _, err := executor.Enqueue(task, opts...); !errors.Is(err, asynq.ErrTaskIDConflict) {
				t.Fatalf("retained task was accepted: %v", err)
			}
			return
		}
		time.Sleep(time.Millisecond)
	}
	t.Fatal("completed task was not retained")
}