xinmi/internal/models/requestbudget/context_test.go

package requestbudget

import (
	"context"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
)

type testGuard struct {
	mu              sync.Mutex
	limit, reserved int
	err             error
	model, phase    string
}

func (g *testGuard) Key() string                     { return "local-test" }
func (g *testGuard) Check(ctx context.Context) error { return ctx.Err() }
func (g *testGuard) Reserve(ctx context.Context, model, phase string) error {
	g.mu.Lock()
	defer g.mu.Unlock()
	if err := ctx.Err(); err != nil {
		return err
	}
	if g.err != nil {
		return g.err
	}
	if g.reserved >= g.limit {
		return ErrStopped
	}
	g.reserved++
	g.model, g.phase = model, phase
	return nil
}

type roundTripFunc func(*http.Request) (*http.Response, error)

func (f roundTripFunc) RoundTrip(r *http.Request) (*http.Response, error) { return f(r) }
func budgetTestContext(g *testGuard) context.Context {
	return WithModel(WithPhase(WithGuard(context.Background(), g), "phase"), "adapter-model")
}

func TestTransportBlocksBeforeBaseOnGuardFailure(t *testing.T) {
	for _, reason := range []string{"budget", "storage error", "missing model", "cancelled"} {
		t.Run(reason, func(t *testing.T) {
			g := &testGuard{limit: 0}
			ctx := budgetTestContext(g)
			if reason == "storage error" {
				g.err = errors.New("durable ledger unavailable")
			}
			if reason == "missing model" {
				ctx = WithGuard(context.Background(), g)
			}
			if reason == "cancelled" {
				var cancel context.CancelFunc
				ctx, cancel = context.WithCancel(ctx)
				cancel()
			}
			sends := 0
			transport := Transport{Base: roundTripFunc(func(*http.Request) (*http.Response, error) { sends++; return nil, errors.New("should not send") })}
			req, _ := http.NewRequestWithContext(ctx, http.MethodPost, "http://local.test", nil)
			_, err := transport.RoundTrip(req)
			if err == nil || sends != 0 || g.reserved != 0 {
				t.Fatalf("failed closed: sends=%d reserved=%d err=%v", sends, g.reserved, err)
			}
		})
	}
}

func TestTransportNeverRefundsFailedSend(t *testing.T) {
	g := &testGuard{limit: 1}
	sends := 0
	transport := Transport{Base: roundTripFunc(func(*http.Request) (*http.Response, error) { sends++; return nil, io.ErrUnexpectedEOF })}
	req, _ := http.NewRequestWithContext(budgetTestContext(g), http.MethodPost, "http://local.test", nil)
	if _, err := transport.RoundTrip(req); !errors.Is(err, io.ErrUnexpectedEOF) {
		t.Fatal(err)
	}
	if _, err := transport.RoundTrip(req); !errors.Is(err, ErrStopped) {
		t.Fatal(err)
	}
	if sends != 1 || g.reserved != 1 || g.model != "adapter-model" || g.phase != "phase" {
		t.Fatalf("reservation lost: %+v sends=%d", g, sends)
	}
}

func TestTransportCountsRedirectsAndPreservesContext(t *testing.T) {
	for _, limit := range []int{1, 2} {
		t.Run(fmt.Sprint(limit), func(t *testing.T) {
			var sends atomic.Int32
			server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				sends.Add(1)
				if r.URL.Path == "/first" {
					http.Redirect(w, r, "/second", http.StatusTemporaryRedirect)
					return
				}
				fmt.Fprint(w, "ok")
			}))
			defer server.Close()
			g := &testGuard{limit: limit}
			client := &http.Client{Transport: Transport{Base: http.DefaultTransport}}
			req, _ := http.NewRequestWithContext(budgetTestContext(g), http.MethodPost, server.URL+"/first", strings.NewReader("body"))
			resp, err := client.Do(req)
			if resp != nil {
				resp.Body.Close()
			}
			if limit == 1 && !errors.Is(err, ErrStopped) {
				t.Fatalf("redirect escaped budget: %v", err)
			}
			if limit == 2 && err != nil {
				t.Fatal(err)
			}
			if int(sends.Load()) != limit || g.reserved != limit || g.model != "adapter-model" || g.phase != "phase" {
				t.Fatalf("redirect context lost: sends=%d guard=%+v", sends.Load(), g)
			}
		})
	}
}

func TestTransportConcurrentSendsCannotExceedGuard(t *testing.T) {
	g := &testGuard{limit: 3}
	var sends atomic.Int32
	transport := Transport{Base: roundTripFunc(func(*http.Request) (*http.Response, error) { sends.Add(1); return nil, io.ErrUnexpectedEOF })}
	var wg sync.WaitGroup
	for i := 0; i < 20; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			req, _ := http.NewRequestWithContext(budgetTestContext(g), http.MethodPost, "http://local.test", nil)
			transport.RoundTrip(req)
		}()
	}
	wg.Wait()
	if sends.Load() != 3 || g.reserved != 3 {
		t.Fatalf("concurrent sends=%d reservations=%d", sends.Load(), g.reserved)
	}
}

func TestTransportPreventsUncountedPostReplay(t *testing.T) {
	for _, controlled := range []bool{false, true} {
		t.Run(fmt.Sprintf("controlled=%t", controlled), func(t *testing.T) {
			var sends atomic.Int32
			server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				io.Copy(io.Discard, r.Body)
				if r.URL.Path == "/bill" && sends.Add(1) == 1 {
					connection, _, err := w.(http.Hijacker).Hijack()
					if err == nil {
						connection.Close()
					}
					return
				}
				fmt.Fprint(w, "ok")
			}))
			defer server.Close()
			base := &http.Transport{}
			defer base.CloseIdleConnections()
			client := &http.Client{Transport: Transport{Base: base}}
			// A reused connection is necessary to exercise Go's transparent retry.
			warm, err := client.Get(server.URL + "/warm")
			if err != nil {
				t.Fatal(err)
			}
			io.Copy(io.Discard, warm.Body)
			warm.Body.Close()
			ctx := context.Background()
			g := &testGuard{limit: 1}
			if controlled {
				ctx = budgetTestContext(g)
			}
			req, _ := http.NewRequestWithContext(ctx, http.MethodPost, server.URL+"/bill", strings.NewReader("body"))
			req.Header.Set("Idempotency-Key", "local-test")
			resp, err := client.Do(req)
			if resp != nil {
				resp.Body.Close()
			}
			if req.GetBody == nil {
				t.Fatal("caller request was mutated")
			}
			if controlled {
				if err == nil || sends.Load() != 1 || g.reserved != 1 {
					t.Fatalf("hidden retry: sends=%d reserved=%d err=%v", sends.Load(), g.reserved, err)
				}
			} else if err != nil || sends.Load() != 2 {
				t.Fatalf("test did not exercise underlying retry: sends=%d err=%v", sends.Load(), err)
			}
		})
	}
}

func TestRequireSupportedOnlyRestrictsControlledContexts(t *testing.T) {
	if err := RequireSupported(context.Background(), false); err != nil {
		t.Fatal(err)
	}
	if err := RequireSupported(budgetTestContext(&testGuard{}), true); err != nil {
		t.Fatal(err)
	}
	if err := RequireSupported(budgetTestContext(&testGuard{}), false); !errors.Is(err, ErrStopped) {
		t.Fatal(err)
	}
}