From 5592ed112799be6e7ee3fb19f01665f7add712f4 Mon Sep 17 00:00:00 2001 From: Marketen Date: Thu, 15 Oct 2026 10:08:41 +0200 Subject: [PATCH] feat: keep the credit hold alive while inference runs Metering expires a credit reservation after its TTL even while the request is still running. Gateway now renews the hold for the life of the request when Metering reports a lease (lease_ms): - renewals at a third of the remaining lease, retried on transient errors no more often than once per second or a tenth of the lease; - the lease is measured on the Gateway's own clock from before the reserve request, so clock skew (an enclave) cannot shorten or extend it; - if Metering refuses a renewal (the hold is gone) or the lease runs out, inference stops with an internal error rather than a user cancellation. Metering versions without lease_ms are unaffected: nothing is renewed. Co-Authored-By: Claude Opus 5.5 --- .../internal/adapters/metering/client.go | 15 +- .../adapters/metering/reservation_lease.go | 88 +++++++++++ .../metering/reservation_lease_test.go | 137 ++++++++++++++++++ .../internal/application/ports/usage.go | 6 + .../internal/application/services/proxy.go | 15 ++ .../services/reservation_keeper_test.go | 63 ++++++++ 6 files changed, 321 insertions(+), 3 deletions(-) create mode 100644 apps/gateway/internal/adapters/metering/reservation_lease.go create mode 100644 apps/gateway/internal/adapters/metering/reservation_lease_test.go create mode 100644 apps/gateway/internal/application/services/reservation_keeper_test.go diff --git a/apps/gateway/internal/adapters/metering/client.go b/apps/gateway/internal/adapters/metering/client.go index a2503c2..00155f8 100644 --- a/apps/gateway/internal/adapters/metering/client.go +++ b/apps/gateway/internal/adapters/metering/client.go @@ -11,6 +11,7 @@ import ( "reflect" "strconv" "strings" + "sync" "time" "github.com/dappnode/dappnode-nexus-gateway/apps/gateway/internal/adapters/observability/metrics" @@ -19,9 +20,10 @@ import ( ) type Client struct { - baseURL string - token string - http *http.Client + baseURL string + token string + http *http.Client + reservations sync.Map // Unclaimed leases by reservation ID, removed by KeepReservation. } func NewClient(baseURL, token string, timeout time.Duration) *Client { @@ -58,12 +60,16 @@ func (c *Client) Reserve(ctx context.Context, auth domain.AuthContext, endpoint MaxOutputTokens: req.MaxOutputTokens, } var res reserveResponse + sent := time.Now() // A lease measured from here never outlasts Metering's. if err := c.postJSON(ctx, "/internal/metering/reservations", payload, &res, 2); err != nil { return "", err } if res.ReservationID == "" { return "", meteringInternalError("reason", "missing reservation_id") } + if res.LeaseMS > 0 { + c.reservations.Store(res.ReservationID, newLease(sent, res.LeaseMS)) + } return res.ReservationID, nil } @@ -365,6 +371,8 @@ func meteringOperation(method, path string) string { return "reserve" case method == http.MethodPost && strings.Contains(requestPath, "/reservations/") && strings.HasSuffix(requestPath, "/complete"): return "complete" + case method == http.MethodPost && strings.Contains(requestPath, "/reservations/") && strings.HasSuffix(requestPath, "/renew"): + return "renew" case method == http.MethodPost && strings.HasSuffix(requestPath, "/usage-failures"): return "record_failure" case method == http.MethodPost && strings.HasSuffix(requestPath, "/authenticate"): @@ -413,6 +421,7 @@ type reserveRequest struct { type reserveResponse struct { ReservationID string `json:"reservation_id"` + LeaseMS int64 `json:"lease_ms,omitempty"` } type completeRequest struct { diff --git a/apps/gateway/internal/adapters/metering/reservation_lease.go b/apps/gateway/internal/adapters/metering/reservation_lease.go new file mode 100644 index 0000000..b0a6e3d --- /dev/null +++ b/apps/gateway/internal/adapters/metering/reservation_lease.go @@ -0,0 +1,88 @@ +package metering + +import ( + "context" + "errors" + "time" + + "github.com/dappnode/dappnode-nexus-gateway/pkg/domain" +) + +// lease is a credit hold measured on this process's clock. Metering sends a +// duration rather than a timestamp, so clock skew (an enclave) cannot make a +// hold look expired or valid when it is not. +type lease struct { + expires time.Time + length time.Duration +} + +func newLease(sent time.Time, ms int64) lease { + length := time.Duration(ms) * time.Millisecond + return lease{expires: sent.Add(length), length: length} +} + +// KeepReservation renews a hold at a third of its remaining time until stop is +// called. Transient failures retry within the lease, never more often than a +// tenth of it or once per second. onLost runs when Metering says the hold is +// gone, or when it expires without a successful renewal. +func (c *Client) KeepReservation(parent context.Context, id string, onLost func()) func() { + value, ok := c.reservations.LoadAndDelete(id) + if !ok { + return func() {} // Older Metering responses do not advertise leases. + } + held := value.(lease) + floor := min(time.Second, held.length/10) + ctx, cancel := context.WithCancel(parent) + done := make(chan struct{}) + lost := func() { + if ctx.Err() == nil { + onLost() + } + } + go func() { + defer close(done) + for { + remaining := time.Until(held.expires) + if remaining <= 0 { + lost() + return + } + timer := time.NewTimer(min(max(remaining/3, floor), remaining)) + select { + case <-ctx.Done(): + timer.Stop() + return + case <-timer.C: + } + if time.Until(held.expires) <= 0 { + continue + } + requestCtx, stop := context.WithDeadline(ctx, held.expires) + var result reserveResponse + sent := time.Now() + err := c.postJSON(requestCtx, "/internal/metering/reservations/"+id+"/renew", struct{}{}, &result, 0) + stop() + switch { + case err == nil && result.ReservationID == id && result.LeaseMS > 0: + if renewed := newLease(sent, result.LeaseMS); renewed.expires.After(held.expires) { + held = renewed + } + case rejectedByMetering(err): + lost() // Completed, expired, or unknown: retrying cannot restore it. + return + } + } + }() + return func() { cancel(); <-done } +} + +// rejectedByMetering reports a 4xx answer; transport errors and 5xx are +// transient and retried while the lease lasts. +func rejectedByMetering(err error) bool { + var gwErr *domain.GatewayError + if !errors.As(err, &gwErr) { + return false + } + status, ok := gwErr.Metadata["metering_status"].(int) + return ok && status >= 400 && status < 500 +} diff --git a/apps/gateway/internal/adapters/metering/reservation_lease_test.go b/apps/gateway/internal/adapters/metering/reservation_lease_test.go new file mode 100644 index 0000000..bd270dd --- /dev/null +++ b/apps/gateway/internal/adapters/metering/reservation_lease_test.go @@ -0,0 +1,137 @@ +package metering + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + "time" + + "github.com/dappnode/dappnode-nexus-gateway/pkg/domain" +) + +func TestReservationRenewsBeyondOriginalExpiryAndStops(t *testing.T) { + var calls atomic.Int32 + var lost atomic.Bool + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if calls.Add(1) == 1 { + w.WriteHeader(503) + return + } + json.NewEncoder(w).Encode(reserveResponse{ReservationID: "lease", LeaseMS: 300}) + })) + defer server.Close() + client := NewClient(server.URL, "fixture", time.Second) + client.reservations.Store("lease", newLease(time.Now(), 300)) + stop := client.KeepReservation(context.Background(), "lease", func() { lost.Store(true) }) + time.Sleep(600 * time.Millisecond) + stop() + if lost.Load() || calls.Load() < 3 { + t.Fatalf("renewal did not preserve hold: lost=%v calls=%d", lost.Load(), calls.Load()) + } + before := calls.Load() + time.Sleep(150 * time.Millisecond) + if calls.Load() != before { + t.Fatal("renewal continued after stop") + } +} + +func TestLostReservationStopsOnlyWhenLeaseExpires(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(503) })) + defer server.Close() + client := NewClient(server.URL, "fixture", time.Second) + client.reservations.Store("lease", newLease(time.Now(), 120)) + lost := make(chan struct{}) + stop := client.KeepReservation(context.Background(), "lease", func() { close(lost) }) + defer stop() + select { + case <-lost: + t.Fatal("temporary error canceled a valid lease") + case <-time.After(50 * time.Millisecond): + } + select { + case <-lost: + case <-time.After(time.Second): + t.Fatal("lost hold did not stop inference") + } +} + +// Metering refusing a renewal means the hold is gone; retrying cannot help. +func TestRejectedRenewalIsLostImmediately(t *testing.T) { + var calls atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls.Add(1) + w.WriteHeader(http.StatusConflict) + json.NewEncoder(w).Encode(map[string]any{"error": map[string]string{"code": "conflict", "message": "reservation is no longer active"}}) + })) + defer server.Close() + client := NewClient(server.URL, "fixture", time.Second) + client.reservations.Store("lease", newLease(time.Now(), 3000)) + lost := make(chan struct{}) + stop := client.KeepReservation(context.Background(), "lease", func() { close(lost) }) + defer stop() + select { + case <-lost: + case <-time.After(1500 * time.Millisecond): + t.Fatal("rejected renewal waited for the lease to run out") + } + if calls.Load() != 1 { + t.Fatalf("retried a rejected renewal: %d calls", calls.Load()) + } +} + +// During a Metering outage each request renews a bounded number of times +// instead of an accelerating burst as the lease runs out. +func TestRenewalRetriesAreSpacedDuringOutage(t *testing.T) { + var calls atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls.Add(1) + w.WriteHeader(503) + })) + defer server.Close() + client := NewClient(server.URL, "fixture", time.Second) + client.reservations.Store("lease", newLease(time.Now(), 1000)) + lost := make(chan struct{}) + stop := client.KeepReservation(context.Background(), "lease", func() { close(lost) }) + defer stop() + select { + case <-lost: + case <-time.After(3 * time.Second): + t.Fatal("lease never expired") + } + // One tenth of a 1 s lease is 100 ms, so at most about ten attempts. + if got := calls.Load(); got > 11 { + t.Fatalf("renewal attempts were not spaced: %d", got) + } +} + +// The lease is measured on the Gateway's clock from before the request, so a +// skewed Metering or enclave clock cannot shorten or extend it. +func TestReserveStoresLeaseRelativeToGatewayClock(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + json.NewEncoder(w).Encode(reserveResponse{ReservationID: "lease", LeaseMS: 60_000}) + })) + defer server.Close() + client := NewClient(server.URL, "fixture", time.Second) + before := time.Now() + id, err := client.Reserve(context.Background(), domain.AuthContext{Account: domain.Account{ID: "a"}, APIKey: domain.APIKey{ID: "k"}}, "chat/completions", domain.GenerateRequest{}, domain.PublicModel{PublicModelID: "m"}, "request") + if err != nil || id != "lease" { + t.Fatalf("reserve = %q, %v", id, err) + } + value, ok := client.reservations.Load("lease") + if !ok { + t.Fatal("lease not stored") + } + held := value.(lease) + if held.length != time.Minute || held.expires.Before(before.Add(time.Minute)) || held.expires.After(time.Now().Add(time.Minute)) { + t.Fatalf("lease = %+v", held) + } +} + +func TestRenewalMetricsDoNotIncludeReservationIDs(t *testing.T) { + if got := meteringOperation(http.MethodPost, "/internal/metering/reservations/private-id/renew"); got != "renew" { + t.Fatalf("renewal metric operation = %q", got) + } +} diff --git a/apps/gateway/internal/application/ports/usage.go b/apps/gateway/internal/application/ports/usage.go index f958530..df62dd0 100644 --- a/apps/gateway/internal/application/ports/usage.go +++ b/apps/gateway/internal/application/ports/usage.go @@ -14,3 +14,9 @@ type UsageMeter interface { RecordSuccess(ctx context.Context, reservationID string, auth domain.AuthContext, endpoint string, req domain.GenerateRequest, resp domain.GenerateResult, model domain.PublicModel, latencyMs int64) error RecordFailure(ctx context.Context, reservationID *string, auth *domain.AuthContext, endpoint string, req *domain.GenerateRequest, model *domain.PublicModel, err error, partialUsage *domain.Usage, latencyMs int64) error } + +// ReservationKeeper maintains a live credit hold while inference progresses. +// onLost stops inference only when renewal cannot preserve the hold. +type ReservationKeeper interface { + KeepReservation(context.Context, string, func()) func() +} diff --git a/apps/gateway/internal/application/services/proxy.go b/apps/gateway/internal/application/services/proxy.go index 9551486..36201d0 100644 --- a/apps/gateway/internal/application/services/proxy.go +++ b/apps/gateway/internal/application/services/proxy.go @@ -13,6 +13,7 @@ import ( "github.com/dappnode/dappnode-nexus-gateway/apps/gateway/internal/adapters/http/middleware" "github.com/dappnode/dappnode-nexus-gateway/apps/gateway/internal/adapters/observability/metrics" + "github.com/dappnode/dappnode-nexus-gateway/apps/gateway/internal/application/ports" "github.com/dappnode/dappnode-nexus-gateway/pkg/domain" "github.com/dappnode/dappnode-nexus-gateway/pkg/observability/logfields" "github.com/google/uuid" @@ -72,7 +73,17 @@ func (s *GenerateService) Proxy(ctx context.Context, endpoint string, raw []byte s.recordTerminalOutcome(metrics.OutcomeError, endpoint, execReq, &model, time.Since(start).Milliseconds()) return nil, err } + ctx, cancelInference := context.WithCancelCause(ctx) + stopRenewal := func() {} + if keeper, ok := s.metering.(ports.ReservationKeeper); ok { + stopRenewal = keeper.KeepReservation(ctx, reservationID, func() { cancelInference(domain.ErrInternal("credit reservation could not be maintained")) }) + } + stopLease := func() { stopRenewal(); cancelInference(nil) } finish := func(target domain.PublicModel, o outcome) { + defer stopLease() + if cause, ok := context.Cause(ctx).(*domain.GatewayError); ok && o.err != nil { + o.err = cause + } o.err = sanitizeErrorWithPIIMapping(o.err, mapping) s.finishProxy(ctx, authCtx, endpoint, execReq, target, requestID, reservationID, start, o) } @@ -138,6 +149,10 @@ func (s *GenerateService) Proxy(ctx context.Context, endpoint string, raw []byte call, err = attempt(executed) } if err != nil { + defer stopLease() + if cause, ok := context.Cause(ctx).(*domain.GatewayError); ok { + err = cause + } err = sanitizeErrorWithPIIMapping(err, mapping) fields := s.buildErrorLogFields(ctx, requestID, &authCtx, endpoint, execReq.PublicModelID, executed, err, time.Since(start).Milliseconds()) s.logger.Error("provider request failed", fields...) diff --git a/apps/gateway/internal/application/services/reservation_keeper_test.go b/apps/gateway/internal/application/services/reservation_keeper_test.go new file mode 100644 index 0000000..c38d864 --- /dev/null +++ b/apps/gateway/internal/application/services/reservation_keeper_test.go @@ -0,0 +1,63 @@ +package services + +import ( + "context" + "strings" + "testing" + + "github.com/dappnode/dappnode-nexus-gateway/apps/gateway/internal/application/ports" + "github.com/dappnode/dappnode-nexus-gateway/pkg/domain" +) + +// keeperMeter is a meter whose credit hold can be lost before inference ends. +type keeperMeter struct { + stubUsageMeter + lose bool + stopped bool +} + +func (m *keeperMeter) KeepReservation(_ context.Context, _ string, onLost func()) func() { + if m.lose { + onLost() + } + return func() { m.stopped = true } +} + +// waitingProvider answers only when its request is canceled, like a slow +// upstream that is still working. +type waitingProvider struct{ stubProvider } + +func (p *waitingProvider) Complete(ctx context.Context, _ []byte, _ domain.PublicModel) ([]byte, *domain.TinfoilTransportProof, error) { + <-ctx.Done() + return nil, nil, ctx.Err() +} + +func keeperService(meter *keeperMeter, provider ports.Provider) *GenerateService { + return NewGenerateService( + &stubAuthService{authCtx: domain.AuthContext{Account: domain.Account{ID: "acc1", Status: domain.AccountStatusActive}, APIKey: domain.APIKey{ID: "key1", Active: true}}}, + &stubModelCatalog{model: domain.PublicModel{PublicModelID: "model", ProviderModelID: "provider-model", ProviderConfig: domain.ProviderConfig{ProviderName: "openai"}, SupportsChatCompletions: true, MaxContextWindow: 1000, MaxOutputTokens: 100}}, + nil, &stubProviderRegistry{providers: map[string]ports.Provider{"openai": provider}}, meter, nil, stubLogger{}, + ) +} + +func TestProxy_StopsInferenceWhenCreditHoldIsLost(t *testing.T) { + meter := &keeperMeter{lose: true} + _, err := keeperService(meter, &waitingProvider{}).Proxy(context.Background(), domain.EndpointChatCompletions, []byte(`{}`), domain.GenerateRequest{PublicModelID: "model"}, "sk-test") + gwErr, ok := err.(*domain.GatewayError) + if !ok || gwErr.Code != domain.ErrCodeInternalError || !strings.Contains(gwErr.Message, "credit reservation") { + t.Fatalf("lost hold reported as %v, want an internal error, not a user cancellation", err) + } + if meter.failureCalls != 1 || !meter.stopped { + t.Fatalf("failure recorded %d times, keeper stopped %v", meter.failureCalls, meter.stopped) + } +} + +func TestProxy_StopsKeepingHoldAfterCompletion(t *testing.T) { + meter := &keeperMeter{} + if _, err := keeperService(meter, &stubProvider{}).Proxy(context.Background(), domain.EndpointChatCompletions, []byte(`{}`), domain.GenerateRequest{PublicModelID: "model"}, "sk-test"); err != nil { + t.Fatal(err) + } + if meter.successCalls != 1 || !meter.stopped { + t.Fatalf("success recorded %d times, keeper stopped %v", meter.successCalls, meter.stopped) + } +}