package pdfpool_test // P5a task 04: the extraction pool. A few slots, a short queue, single-flight per key, and at // most one job in flight per caller. import ( "context" "errors" "sync" "sync/atomic" "testing" "time" "git.wntrmute.dev/kyle/sift/internal/pdfpool" ) func waitUntil(t *testing.T, cond func() bool) { t.Helper() deadline := time.Now().Add(3 * time.Second) for !cond() { if time.Now().After(deadline) { t.Fatal("condition not reached within 3 s") } time.Sleep(5 * time.Millisecond) } } func TestSingleFlight(t *testing.T) { p := pdfpool.New(2, 4, time.Second) var calls atomic.Int32 release := make(chan struct{}) fn := func(ctx context.Context) (any, error) { calls.Add(1) <-release return "text", nil } var wg sync.WaitGroup results := make(chan any, 5) for i := 0; i < 5; i++ { wg.Add(1) go func(i int) { defer wg.Done() v, err := p.Do(context.Background(), "doc-1", "caller-"+string(rune('a'+i)), fn) if err != nil { t.Error(err) } results <- v }(i) } waitUntil(t, func() bool { return calls.Load() == 1 && p.Waiting("doc-1") == 4 }) close(release) wg.Wait() close(results) for v := range results { if v != "text" { t.Errorf("a sharer got %v", v) } } if n := calls.Load(); n != 1 { t.Errorf("fn ran %d times for one key, want 1", n) } // Nothing is cached: the next call runs fn again. _, _ = p.Do(context.Background(), "doc-1", "x", func(ctx context.Context) (any, error) { calls.Add(1); return nil, nil }) if n := calls.Load(); n != 2 { t.Errorf("after the flight ended, fn ran %d times in total, want 2", n) } } // Two slots and four waiting places: the seventh distinct job is refused at once. func TestSlotsAndQueue(t *testing.T) { p := pdfpool.New(2, 4, 5*time.Second) release := make(chan struct{}) var running atomic.Int32 block := func(ctx context.Context) (any, error) { running.Add(1) defer running.Add(-1) <-release return nil, nil } var wg sync.WaitGroup for i := 0; i < 6; i++ { wg.Add(1) go func(i int) { defer wg.Done() if _, err := p.Do(context.Background(), "doc-"+string(rune('a'+i)), "caller-"+string(rune('a'+i)), block); err != nil { t.Errorf("job %d: %v", i, err) } }(i) } waitUntil(t, func() bool { return running.Load() == 2 && p.Queued() == 4 }) start := time.Now() _, err := p.Do(context.Background(), "doc-z", "caller-z", block) if !errors.Is(err, pdfpool.ErrBusy) { t.Errorf("seventh job: %v, want ErrBusy", err) } if d := time.Since(start); d > 200*time.Millisecond { t.Errorf("a full queue refused after %v; it must refuse at once", d) } if running.Load() > 2 { t.Errorf("%d jobs ran at once, want at most 2", running.Load()) } close(release) wg.Wait() } // A queued job that waits longer than Wait is refused with ErrBusy and gives up its place. func TestQueueWait(t *testing.T) { p := pdfpool.New(1, 1, 150*time.Millisecond) release := make(chan struct{}) defer close(release) go func() { _, _ = p.Do(context.Background(), "a", "ca", func(ctx context.Context) (any, error) { <-release; return nil, nil }) }() waitUntil(t, func() bool { return p.Running() == 1 }) start := time.Now() _, err := p.Do(context.Background(), "b", "cb", func(ctx context.Context) (any, error) { return nil, nil }) if !errors.Is(err, pdfpool.ErrBusy) { t.Fatalf("waiter: %v, want ErrBusy", err) } if d := time.Since(start); d < 100*time.Millisecond || d > time.Second { t.Errorf("waiter refused after %v, want about 150 ms", d) } if p.Queued() != 0 { t.Errorf("queued = %d after the waiter gave up, want 0", p.Queued()) } } func TestCallerBusy(t *testing.T) { p := pdfpool.New(2, 4, time.Second) release := make(chan struct{}) go func() { _, _ = p.Do(context.Background(), "doc-1", "straylight", func(ctx context.Context) (any, error) { <-release; return nil, nil }) }() waitUntil(t, func() bool { return p.Running() == 1 }) if _, err := p.Do(context.Background(), "doc-2", "straylight", func(ctx context.Context) (any, error) { return nil, nil }); !errors.Is(err, pdfpool.ErrCallerBusy) { t.Errorf("second job from the same caller: %v, want ErrCallerBusy", err) } // The same caller asking for the same document again joins the flight instead. done := make(chan error, 1) go func() { _, err := p.Do(context.Background(), "doc-1", "straylight", func(ctx context.Context) (any, error) { return nil, nil }) done <- err }() waitUntil(t, func() bool { return p.Waiting("doc-1") == 1 }) close(release) if err := <-done; err != nil { t.Errorf("joining its own flight: %v", err) } // Once the job is done the caller may start another. if _, err := p.Do(context.Background(), "doc-2", "straylight", func(ctx context.Context) (any, error) { return nil, nil }); err != nil { t.Errorf("after the first job ended: %v", err) } } func TestErrorsAndPanicsReachEverySharer(t *testing.T) { p := pdfpool.New(2, 4, time.Second) boom := errors.New("boom") if _, err := p.Do(context.Background(), "k", "c", func(ctx context.Context) (any, error) { return nil, boom }); !errors.Is(err, boom) { t.Errorf("fn error: %v", err) } _, err := p.Do(context.Background(), "k", "c", func(ctx context.Context) (any, error) { panic("kaboom") }) if err == nil { t.Fatal("a panicking fn returned no error") } if p.Running() != 0 { t.Errorf("running = %d after a panic, want 0 (the slot leaked)", p.Running()) } if _, err := p.Do(context.Background(), "k2", "c", func(ctx context.Context) (any, error) { return 1, nil }); err != nil { t.Errorf("the caller is still marked busy after its job panicked: %v", err) } } func TestContextCancelWhileQueued(t *testing.T) { p := pdfpool.New(1, 2, 5*time.Second) release := make(chan struct{}) defer close(release) go func() { _, _ = p.Do(context.Background(), "a", "ca", func(ctx context.Context) (any, error) { <-release; return nil, nil }) }() waitUntil(t, func() bool { return p.Running() == 1 }) ctx, cancel := context.WithCancel(context.Background()) errc := make(chan error, 1) go func() { _, err := p.Do(ctx, "b", "cb", func(ctx context.Context) (any, error) { return nil, nil }) errc <- err }() waitUntil(t, func() bool { return p.Queued() == 1 }) cancel() if err := <-errc; !errors.Is(err, context.Canceled) { t.Errorf("cancelled waiter: %v, want context.Canceled", err) } if p.Queued() != 0 { t.Errorf("queued = %d after cancel, want 0", p.Queued()) } } // The job's own context is not the first caller's: if the caller that started a shared flight // goes away, the sharers still get the result. func TestFlightOutlivesItsStarter(t *testing.T) { p := pdfpool.New(1, 2, 5*time.Second) release := make(chan struct{}) ctx, cancel := context.WithCancel(context.Background()) jobCtxErr := make(chan error, 1) go func() { _, _ = p.Do(ctx, "doc", "first", func(jctx context.Context) (any, error) { <-release jobCtxErr <- jctx.Err() return "done", nil }) }() waitUntil(t, func() bool { return p.Running() == 1 }) res := make(chan any, 1) go func() { v, _ := p.Do(context.Background(), "doc", "second", nil) res <- v }() waitUntil(t, func() bool { return p.Waiting("doc") == 1 }) cancel() close(release) if v := <-res; v != "done" { t.Errorf("sharer got %v after the starter left, want done", v) } if e := <-jobCtxErr; e != nil { t.Errorf("the job's context was cancelled with its starter: %v", e) } }