Something went wrong. Try again.
An ATProto Powered Reputation Network
Something went wrong. Try again.
11 kB · 432 lines
Go
at master
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433package upstream
import ( "context" "errors" "io" "log/slog" "math/rand/v2" "net" "net/http" "strconv" "strings" "sync" "sync/atomic" "syscall" "time"
"provenance/internal/logging")
const ( UserAgent = "provenance/0.1 (https://provenance.ngp.computer)" DefaultWorkers = 4 MaxWorkers = 64 MaxAttempts = 3 queuePerWorker = 8 maxCooldowns = 1024 retryBase = time.Second retryMax = 30 * time.Second retryDrainBytes = 4 << 10 unknownBodyLength = -1)
type job struct { req *http.Request base http.RoundTripper result chan result policy RetryPolicy queuedAt time.Time}
type result struct { resp *http.Response err error}
// Queue bounds external HTTP work, including response reads and retry waits.type Queue struct { ctx context.Context cancel context.CancelFunc jobs chan job workers sync.WaitGroup mu sync.Mutex cooldowns map[string]time.Time now func() time.Time wait func(context.Context, time.Duration) error}
func New(ctx context.Context, workers int) (*Queue, error) { if workers < 1 || workers > MaxWorkers { return nil, errors.New("external workers must be between 1 and 64") } ctx, cancel := context.WithCancel(ctx) q := &Queue{ctx: ctx, cancel: cancel, jobs: make(chan job, workers*queuePerWorker), cooldowns: make(map[string]time.Time), now: time.Now, wait: Wait} for range workers { q.workers.Go(q.run) } return q, nil}
func (q *Queue) Close() { q.cancel(); q.workers.Wait() }
var defaultQueue = sync.OnceValue(func() *Queue { q, _ := New(context.Background(), DefaultWorkers); return q })
func Default() *Queue { return defaultQueue() }
type queued struct { queue *Queue base http.RoundTripper policy RetryPolicy}
type RetryPolicy int
const ( QueueRetries RetryPolicy = iota CallerRetries)
func (q *Queue) Transport(base http.RoundTripper) http.RoundTripper { return q.TransportWithPolicy(base, QueueRetries)}
// An archive decoder owns range retries. Do not multiply its attempt budget in the queue.func (q *Queue) TransportWithPolicy(base http.RoundTripper, policy RetryPolicy) http.RoundTripper { if base == nil { base = http.DefaultTransport } return queued{q, base, policy}}
func (t queued) RoundTrip(req *http.Request) (resp *http.Response, err error) { defer func() { if req.Body != nil && err != nil { req.Body.Close() } }() if err := req.Context().Err(); err != nil { return nil, err } if err := t.queue.ctx.Err(); err != nil { return nil, err } work := job{req: req, base: t.base, result: make(chan result), policy: t.policy, queuedAt: time.Now()} select { case <-req.Context().Done(): return nil, req.Context().Err() case <-t.queue.ctx.Done(): return nil, t.queue.ctx.Err() case t.queue.jobs <- work: } select { case <-req.Context().Done(): return nil, req.Context().Err() case <-t.queue.ctx.Done(): return nil, t.queue.ctx.Err() case got := <-work.result: return got.resp, got.err }}
func (q *Queue) run() { for { if q.ctx.Err() != nil { return } select { case <-q.ctx.Done(): return case work := <-q.jobs: q.serve(work) } }}
func (q *Queue) serve(work job) { ctx, cancel := context.WithCancel(work.req.Context()) stop := context.AfterFunc(q.ctx, cancel) if q.ctx.Err() != nil { cancel() } cleanup := func() { stop(); cancel() } req := work.req.Clone(ctx) req.Header.Set("User-Agent", UserAgent) started := time.Now() logger := slog.Default().With("component", "external_http", "method", req.Method, "host", req.URL.Host, "path", req.URL.Path, "retry_owner", retryOwner(work.policy)) logger.Debug("upstream request started", "queue_wait", started.Sub(work.queuedAt)) resp, err := q.fetch(req, work.base, work.policy) if err != nil { level := slog.LevelWarn if errors.Is(err, context.Canceled) { level = slog.LevelDebug } logger.Log(req.Context(), level, "upstream request failed", "duration", time.Since(started), "error", logging.Error(err)) } var body *heldBody if resp != nil { logger = logger.With("status", resp.StatusCode) logger.Debug("upstream response received", "header_duration", time.Since(started), "content_length", resp.ContentLength) if resp.StatusCode >= http.StatusBadRequest { logger.Warn("upstream response failed") } body = &heldBody{ReadCloser: resp.Body, done: make(chan struct{}), cleanup: cleanup, remaining: resp.ContentLength} body.logger = logger body.onClose = func() { logger.Debug("upstream response closed", "duration", time.Since(started), "read_bytes", body.readBytes.Load()) } if resp.StatusCode == http.StatusSwitchingProtocols { body.remaining = unknownBodyLength } resp.Body = body if resp.StatusCode != http.StatusSwitchingProtocols { defer body.Close() } } else { defer cleanup() } select { case work.result <- result{resp, err}: case <-ctx.Done(): if body != nil { body.Close() } return } if body == nil { return } // A WebSocket owns its connection after the handshake. Do not occupy a fetch worker. if resp.StatusCode == http.StatusSwitchingProtocols { return } if resp.ContentLength == 0 { body.Close() return } select { case <-body.done: case <-ctx.Done(): }}
type heldBody struct { io.ReadCloser done chan struct{} once sync.Once cleanup func() remaining int64 readBytes atomic.Int64 logger *slog.Logger onClose func()}
func (b *heldBody) Close() error { var err error b.once.Do(func() { err = b.ReadCloser.Close() b.cleanup() close(b.done) if b.onClose != nil { b.onClose() } if err != nil && b.logger != nil { b.logger.Warn("upstream response close failed", "error", logging.Error(err)) } }) return err}
func (b *heldBody) Write(p []byte) (int, error) { writer, ok := b.ReadCloser.(io.Writer) if !ok { return 0, errors.New("upstream response is not writable") } return writer.Write(p)}
func (b *heldBody) Read(p []byte) (int, error) { if b.remaining == 0 { return 0, io.EOF } n, err := b.ReadCloser.Read(p) b.readBytes.Add(int64(n)) if err != nil && !errors.Is(err, io.EOF) && b.logger != nil { b.logger.Warn("upstream response read failed", "read_bytes", b.readBytes.Load(), "error", logging.Error(err)) } // ReadFull can consume a complete range without a final EOF read. Release its worker at the declared length. if b.remaining > 0 { b.remaining -= int64(n) } if err != nil || b.remaining == 0 { b.Close() } return n, err}
func origin(req *http.Request) string { port := req.URL.Port() if port == "" { port = "443" if req.URL.Scheme == "http" { port = "80" } } return req.URL.Scheme + "://" + net.JoinHostPort(strings.ToLower(req.URL.Hostname()), port)}
func (q *Queue) fetch(req *http.Request, base http.RoundTripper, policy RetryPolicy) (*http.Response, error) { host := origin(req) safe := (req.Method == http.MethodGet || req.Method == http.MethodHead) && (req.Body == nil || req.Body == http.NoBody) delay := retryBase for attempt := 1; ; attempt++ { if err := q.cool(req.Context(), host); err != nil { return nil, err } resp, err := base.RoundTrip(req) if err != nil && resp != nil { resp.Body.Close() resp = nil } if req.Context().Err() != nil { if resp != nil { resp.Body.Close() } return nil, req.Context().Err() } retry := transientError(err) if retry { if cooldownErr := q.deferHost(host, Jitter(delay)); cooldownErr != nil { return nil, cooldownErr } } if resp != nil { retry = transient(resp.StatusCode) if retry { pause := max(retryAfter(resp.Header.Get("Retry-After"), q.now()), Jitter(delay)) if err := q.deferHost(host, pause); err != nil { resp.Body.Close() return nil, err } } } if !retry || !safe || attempt >= MaxAttempts || policy == CallerRetries { return resp, err } if resp != nil { // Drain small error responses to reuse connections without reading unbounded bodies. _, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, retryDrainBytes)) resp.Body.Close() } status := 0 if resp != nil { status = resp.StatusCode } slog.Warn("upstream retry scheduled", "component", "external_http", "method", req.Method, "host", req.URL.Host, "path", req.URL.Path, "status", status, "attempt", attempt, "next_attempt", attempt+1, "backoff", delay, "error", logging.Error(err)) delay = min(delay*2, retryMax) }}
func retryOwner(policy RetryPolicy) string { if policy == CallerRetries { return "caller" } return "queue"}
func transient(status int) bool { return status == http.StatusTooManyRequests || status == http.StatusBadGateway || status == http.StatusServiceUnavailable || status == http.StatusGatewayTimeout || status == http.StatusInternalServerError}
func transientError(err error) bool { if err == nil || errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { return false } if errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) || errors.Is(err, syscall.ECONNRESET) || errors.Is(err, syscall.ECONNREFUSED) || errors.Is(err, syscall.EPIPE) { return true } var network net.Error return errors.As(err, &network) && (network.Timeout() || network.Temporary())}
func (q *Queue) deferHost(host string, delay time.Duration) error { q.mu.Lock() defer q.mu.Unlock() now := q.now() if len(q.cooldowns) >= maxCooldowns { for key, until := range q.cooldowns { if !until.After(now) { delete(q.cooldowns, key) } } if _, exists := q.cooldowns[host]; !exists && len(q.cooldowns) >= maxCooldowns { return errors.New("external cooldown capacity exceeded") } } q.cooldowns[host] = maxTime(q.cooldowns[host], now.Add(delay)) return nil}
func maxTime(a, b time.Time) time.Time { if a.After(b) { return a } return b}
func (q *Queue) cool(ctx context.Context, host string) error { for { if err := ctx.Err(); err != nil { return err } q.mu.Lock() until := q.cooldowns[host] now := q.now() if !until.After(now) { delete(q.cooldowns, host) } q.mu.Unlock() if !until.After(now) { return nil } slog.Debug("upstream cooldown wait", "component", "external_http", "origin", host, "retry_after", until.Sub(now)) if err := q.wait(ctx, until.Sub(now)); err != nil { return err } }}
func retryAfter(raw string, now time.Time) time.Duration { if seconds, err := strconv.ParseUint(raw, 10, 64); err == nil || errors.Is(err, strconv.ErrRange) { // Reject overflow without converting an upstream cooldown into an immediate retry. const maxSeconds = uint64((1<<63 - 1) / int64(time.Second)) if err != nil || seconds > maxSeconds { return time.Duration(1<<63 - 1) } return time.Duration(seconds) * time.Second } if date, err := http.ParseTime(raw); err == nil { return max(0, date.Sub(now)) } return 0}
// Jitter preserves a minimum delay and spreads retries across the latter half of the interval.func Jitter(delay time.Duration) time.Duration { if delay <= 1 { return delay } half := delay / 2 return half + time.Duration(rand.Int64N(int64(delay-half)+1))}
func Wait(ctx context.Context, delay time.Duration) error { timer := time.NewTimer(delay) defer timer.Stop() select { case <-ctx.Done(): return ctx.Err() case <-timer.C: return nil }}