fix(gateway): stop leaking quota reservations and rule keys

internal/gateway/engine.go

@@ -6,6 +6,7 @@ package gateway
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"sync"
 )
@@ -112,8 +113,12 @@ func (e *Engine) Stop(ctx context.Context) error {
 
 	var firstErr error
 	for key := range e.active {
-		if err := e.removeRuleLocked(ctx, key); err != nil && firstErr == nil {
-			firstErr = fmt.Errorf("stop: remove rule %s: %w", key, err)
+		// A key whose teardown failed stays active whether or not an
+		// earlier key already failed; only the first error is returned.
+		if err := e.removeRuleLocked(ctx, key); err != nil {
+			if firstErr == nil {
+				firstErr = fmt.Errorf("stop: remove rule %s: %w", key, err)
+			}
 			continue
 		}
 		delete(e.active, key)
@@ -143,6 +148,14 @@ func (e *Engine) applyRuleLocked(ctx context.Context, rule DesiredRule) error {
 	}
 
 	if err := e.datapath.ApplyRule(ctx, rule); err != nil {
+		// Reconcile only adds the key to e.active on success, so no later
+		// removeRuleLocked would ever release the reservation made above.
+		if relErr := e.quota.Release(ctx, rule.Key); relErr != nil {
+			return errors.Join(
+				fmt.Errorf("apply datapath rule %s: %w", rule.Key, err),
+				fmt.Errorf("release quota for %s: %w", rule.Key, relErr),
+			)
+		}
 		return fmt.Errorf("apply datapath rule %s: %w", rule.Key, err)
 	}
 
@@ -153,11 +166,18 @@ func (e *Engine) applyRuleLocked(ctx context.Context, rule DesiredRule) error {
 // removeRuleLocked tears down a single rule's datapath and quota state.
 // Caller must hold e.mu.
 func (e *Engine) removeRuleLocked(ctx context.Context, key string) error {
-	if err := e.datapath.RemoveRule(ctx, key); err != nil {
-		return fmt.Errorf("remove datapath rule %s: %w", key, err)
+	// Quota is released even when datapath removal fails, so a reservation
+	// cannot outlive the rule; Release is a no-op for an already-released
+	// key, so the caller's next retry stays correct.
+	dpErr := e.datapath.RemoveRule(ctx, key)
+	if dpErr != nil {
+		dpErr = fmt.Errorf("remove datapath rule %s: %w", key, dpErr)
 	}
 	if err := e.quota.Release(ctx, key); err != nil {
-		return fmt.Errorf("release quota for %s: %w", key, err)
+		return errors.Join(dpErr, fmt.Errorf("release quota for %s: %w", key, err))
+	}
+	if dpErr != nil {
+		return dpErr
 	}
 
 	e.telemetry.RuleRemoved(ctx, key)

internal/gateway/engine_test.go

@@ -7,6 +7,7 @@ package gateway
 import (
 	"context"
 	"errors"
+	"strings"
 	"testing"
 )
 
@@ -187,6 +188,59 @@ func TestEngine_ReconcileApplyErrorReportsUnhealthyKeepsGoing(t *testing.T) {
 	}
 }
 
+func TestEngine_ReconcileApplyErrorReleasesQuotaReservation(t *testing.T) {
+	dp := newFakeDatapath()
+	dp.applyErr = errors.New("simulated datapath failure")
+	quota := &fakeQuota{}
+	e := NewEngine(dp, quota, &fakeTelemetry{})
+
+	desired := EngineState{Rules: map[string]DesiredRule{testKeyA: {Key: testKeyA}}}
+	if _, err := e.Reconcile(context.Background(), desired); err != nil {
+		t.Fatalf("Reconcile: %v", err)
+	}
+
+	// Reconcile never adds a failed key to e.active, so this is the only
+	// chance the reservation ever gets to be released.
+	if len(quota.released) != 1 || quota.released[0] != testKeyA {
+		t.Errorf("quota.released = %v, want [%s] after a failed apply", quota.released, testKeyA)
+	}
+}
+
+func TestEngine_ReconcileRemoveErrorKeepsRuleActiveAndReleasesQuota(t *testing.T) {
+	dp := newFakeDatapath()
+	quota := &fakeQuota{}
+	e := NewEngine(dp, quota, &fakeTelemetry{})
+	ctx := context.Background()
+
+	if _, err := e.Reconcile(ctx, EngineState{Rules: map[string]DesiredRule{testKeyA: {Key: testKeyA}}}); err != nil {
+		t.Fatalf("first Reconcile: %v", err)
+	}
+
+	dp.removeErr = errors.New("simulated teardown failure")
+	status, err := e.Reconcile(ctx, EngineState{Rules: map[string]DesiredRule{}})
+	if err != nil {
+		t.Fatalf("second Reconcile: %v", err)
+	}
+	if status.Healthy {
+		t.Error("status.Healthy = true, want false (datapath teardown failed)")
+	}
+	if len(status.Rules) != 1 || !strings.Contains(status.Rules[0].Error, "simulated teardown failure") {
+		t.Errorf("status.Rules = %+v, want one entry reporting the datapath error", status.Rules)
+	}
+	if len(quota.released) != 1 || quota.released[0] != testKeyA {
+		t.Errorf("quota.released = %v, want [%s] even though datapath removal failed", quota.released, testKeyA)
+	}
+
+	// The key stays active so the next pass retries the datapath teardown.
+	active, err := e.Status(ctx)
+	if err != nil {
+		t.Fatalf("Status: %v", err)
+	}
+	if len(active.Rules) != 1 || active.Rules[0].Key != testKeyA {
+		t.Errorf("Status() = %+v, want %q still active after a failed teardown", active, testKeyA)
+	}
+}
+
 func TestEngine_StatusReflectsActiveRulesWithoutReconciling(t *testing.T) {
 	dp := newFakeDatapath()
 	e := NewEngine(dp, &fakeQuota{}, &fakeTelemetry{})
@@ -231,6 +285,33 @@ func TestEngine_StopTearsDownEveryActiveRule(t *testing.T) {
 	}
 }
 
+func TestEngine_StopKeepsEveryFailedTeardownActive(t *testing.T) {
+	dp := newFakeDatapath()
+	e := NewEngine(dp, &fakeQuota{}, &fakeTelemetry{})
+	ctx := context.Background()
+
+	desired := EngineState{Rules: map[string]DesiredRule{testKeyA: {Key: testKeyA}, testKeyB: {Key: testKeyB}}}
+	if _, err := e.Reconcile(ctx, desired); err != nil {
+		t.Fatalf("Reconcile: %v", err)
+	}
+
+	dp.removeErr = errors.New("simulated teardown failure")
+	if err := e.Stop(ctx); err == nil {
+		t.Fatal("Stop() = nil, want the first teardown error")
+	}
+
+	// Both teardowns failed, so neither may be dropped -- whichever key
+	// map iteration reaches second must not fall through to the delete
+	// just because the first one already set firstErr.
+	status, err := e.Status(ctx)
+	if err != nil {
+		t.Fatalf("Status: %v", err)
+	}
+	if len(status.Rules) != 2 {
+		t.Errorf("Status() after a wholly-failed Stop = %+v, want both rules still active", status)
+	}
+}
+
 func TestEngine_DatapathGenerationDelegates(t *testing.T) {
 	dp := newFakeDatapath()
 	dp.generation = 12345

internal/gateway/kerneldatapath.go

@@ -8,6 +8,7 @@ import (
 	"context"
 	"fmt"
 	"net/netip"
+	"slices"
 	"sync"
 
 	"go.datum.net/galactic/internal/plumbing/ebpf/edgemap"
@@ -107,8 +108,16 @@ func (d *KernelDatapath) ApplyRule(_ context.Context, rule DesiredRule) error {
 		backends[i] = edgemap.Backend{Addr: b.Address, Port: b.Port, USID: b.USID}
 	}
 
-	for _, key := range keys {
+	for i, key := range keys {
 		if err := d.ruleTable.Register(key, backends); err != nil {
+			// Record the keys that did land, alongside the ones this rule
+			// already owned (the prune below has not run yet), so
+			// RemoveRule can still find every live entry.
+			for _, written := range keys[:i] {
+				if !slices.Contains(d.ruleKeysByName[rule.Key], written) {
+					d.ruleKeysByName[rule.Key] = append(d.ruleKeysByName[rule.Key], written)
+				}
+			}
 			return fmt.Errorf("kerneldatapath: apply rule %s: %w", rule.Key, err)
 		}
 	}