feat: Distribute private DNS configuration across regional fleets

internal/internaldns/platform/allocator.go

@@ -10,6 +10,7 @@ import (
 	"net/netip"
 	"time"
 
+	dnsv1alpha1 "go.miloapis.com/dns-operator/api/v1alpha1"
 	"go.miloapis.com/dns-operator/internal/internaldns/model"
 	corev1 "k8s.io/api/core/v1"
 	apierrors "k8s.io/apimachinery/pkg/api/errors"
@@ -46,9 +47,9 @@ func (a *Allocator) Allocate(ctx context.Context, projectUID, contextUID types.U
 // destination across every project sharing this platform control plane. Claims
 // are retained after revocation so a stale route cannot become valid for a new
 // context lifetime.
-func (a *Allocator) ClaimDestination(ctx context.Context, projectUID, accessUID types.UID, region, address string, port int32) error {
+func (a *Allocator) ClaimDestination(ctx context.Context, projectUID, contextUID types.UID, region, address string, port int32) error {
 	parsed, err := netip.ParseAddr(address)
-	if err != nil || !parsed.IsGlobalUnicast() || projectUID == "" || accessUID == "" || region == "" || port != 53 {
+	if err != nil || !parsed.IsGlobalUnicast() || projectUID == "" || contextUID == "" || region == "" || port != 53 {
 		return fmt.Errorf("destination claim identity, region, IP, and DNS port are required")
 	}
 	parsed = parsed.Unmap()
@@ -67,14 +68,29 @@ func (a *Allocator) ClaimDestination(ctx context.Context, projectUID, accessUID
 	} else if !apierrors.IsNotFound(err) {
 		return err
 	}
-	owner := string(projectUID) + "/" + string(accessUID)
+	owner := string(projectUID) + "/" + string(contextUID)
 	claim := &corev1.ConfigMap{ObjectMeta: metav1.ObjectMeta{Name: "dns-destination-claim-" + model.OpaqueToken(region+"/"+address+"/"+fmt.Sprint(port)), Namespace: a.Namespace, Labels: map[string]string{leaseLabel: "true"}}, Data: map[string]string{"owner": owner, "region": region, "address": address, "port": fmt.Sprint(port)}}
 	if err := a.Client.Create(ctx, claim); apierrors.IsAlreadyExists(err) {
 		var existing corev1.ConfigMap
 		if err := a.Client.Get(ctx, client.ObjectKeyFromObject(claim), &existing); err != nil {
 			return err
 		}
 		if existing.Data["owner"] != owner || existing.Data["region"] != region || existing.Data["address"] != address || existing.Data["port"] != fmt.Sprint(port) {
+			var bindings dnsv1alpha1.DNSResolverBindingList
+			if err := a.Client.List(ctx, &bindings, client.InNamespace(a.Namespace)); err != nil {
+				return err
+			}
+			for _, binding := range bindings.Items {
+				b := binding.Spec
+				if b.Source.ProjectUID == projectUID && b.Source.ResolverContextRef.UID == contextUID &&
+					b.Source.AccessBindingRef.UID != "" && existing.Data["owner"] == string(projectUID)+"/"+string(b.Source.AccessBindingRef.UID) &&
+					b.Placement.Region == region && b.Configuration.Query.NodeAddress == address && port == 53 &&
+					existing.Data["region"] == region && existing.Data["address"] == address && existing.Data["port"] == fmt.Sprint(port) {
+					base := existing.DeepCopy()
+					existing.Data["owner"] = owner
+					return a.Client.Patch(ctx, &existing, client.MergeFromWithOptions(base, client.MergeFromWithOptimisticLock{}))
+				}
+			}
 			return fmt.Errorf("service destination is already reserved")
 		}
 		return nil

internal/internaldns/platform/allocator_test.go

@@ -77,6 +77,51 @@ func TestAddressClaimsAreConcurrentStableAndLifetimeScoped(t *testing.T) {
 	}
 }
 
+func TestDestinationClaimBelongsToContextLifetime(t *testing.T) {
+	ctx := context.Background()
+	cl := testClient(t)
+	allocator := Allocator{Client: cl, Namespace: "platform"}
+	for range 2 {
+		if err := allocator.ClaimDestination(ctx, "project", "context", "east", "fd70::53", 53); err != nil {
+			t.Fatal(err)
+		}
+	}
+	if err := allocator.ClaimDestination(ctx, "project", "replacement-context", "east", "fd70::53", 53); err == nil {
+		t.Fatal("a replacement context reused a quarantined destination")
+	}
+	if err := allocator.ClaimDestination(ctx, "another-project", "context", "east", "fd70::53", 53); err == nil {
+		t.Fatal("another project reused a destination")
+	}
+}
+
+func TestDestinationClaimConversionRequiresPinnedBinding(t *testing.T) {
+	ctx := context.Background()
+	cl := testClient(t)
+	allocator := Allocator{Client: cl, Namespace: "platform"}
+	if err := allocator.ClaimDestination(ctx, "project", "access", "east", "fd70::53", 53); err != nil {
+		t.Fatal(err)
+	}
+	if err := allocator.ClaimDestination(ctx, "project", "context", "east", "fd70::53", 53); err == nil {
+		t.Fatal("claim changed ownership without lifetime proof")
+	}
+	binding := &dnsv1alpha1.DNSResolverBinding{}
+	binding.Namespace, binding.Name = "platform", "binding"
+	binding.Spec.Source.ProjectUID = "project"
+	binding.Spec.Source.ResolverContextRef.UID = "context"
+	binding.Spec.Source.AccessBindingRef.UID = "access"
+	binding.Spec.Placement.Region = "east"
+	binding.Spec.Configuration.Query.NodeAddress = "fd70::53"
+	if err := cl.Create(ctx, binding); err != nil {
+		t.Fatal(err)
+	}
+	if err := allocator.ClaimDestination(ctx, "project", "context", "east", "fd70::53", 53); err != nil {
+		t.Fatal(err)
+	}
+	if err := allocator.ClaimDestination(ctx, "project", "access", "east", "fd70::53", 53); err == nil {
+		t.Fatal("converted claim returned to access lifetime")
+	}
+}
+
 func TestServiceDestinationsCannotOverlapRegionalAddressPool(t *testing.T) {
 	ctx := context.Background()
 	for _, prefix := range []string{"", "fd80::/96"} {

internal/internaldns/platform/capacity.go

@@ -0,0 +1,113 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+
+package platform
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"sort"
+	"time"
+
+	dnsv1alpha1 "go.miloapis.com/dns-operator/api/v1alpha1"
+	"go.miloapis.com/dns-operator/internal/internaldns/model"
+	"sigs.k8s.io/controller-runtime/pkg/client"
+)
+
+func fitSnapshot(snapshot model.ServingSnapshot, prior model.ServingSnapshot, maxBindings, maxBytes int) (model.ServingSnapshot, map[string]bool, error) {
+	previous := map[string]model.Binding{}
+	for _, binding := range prior.Bindings {
+		previous[binding.BindingUID] = binding
+	}
+	candidates := append([]model.Binding(nil), snapshot.Bindings...)
+	sort.Slice(candidates, func(i, j int) bool {
+		_, first := previous[candidates[i].BindingUID]
+		_, second := previous[candidates[j].BindingUID]
+		if first != second {
+			return first
+		}
+		return candidates[i].BindingUID < candidates[j].BindingUID
+	})
+	snapshot.Bindings = []model.Binding{}
+	rejected := map[string]bool{}
+	envelope, err := model.NewEnvelope(model.KindServingSnapshot, "dns-serving-"+model.OpaqueToken(snapshot.Region+"/"+snapshot.Shard)+"-e18446744073709551615-r18446744073709551615", snapshot.Region, snapshot.Shard, "00000000-0000-0000-0000-000000000000", snapshot.ConfigurationEpoch, snapshot.ConfigurationRevision, time.Now(), snapshot)
+	if err != nil {
+		return snapshot, rejected, err
+	}
+	payload, err := json.Marshal(envelope)
+	if err != nil {
+		return snapshot, rejected, err
+	}
+	used := len(payload)
+	fits := func(binding model.Binding) bool {
+		if len(snapshot.Bindings) >= maxBindings {
+			return false
+		}
+		data, err := json.Marshal(binding)
+		if err != nil {
+			return false
+		}
+		size := len(data)
+		if len(snapshot.Bindings) > 0 {
+			size++
+		}
+		if used+size > maxBytes {
+			return false
+		}
+		used += size
+		return true
+	}
+	for _, binding := range candidates {
+		if fits(binding) {
+			snapshot.Bindings = append(snapshot.Bindings, binding)
+			continue
+		}
+		rejected[binding.BindingUID] = true
+		if old, ok := previous[binding.BindingUID]; ok {
+			if safeExpansion(old, binding) && fits(old) {
+				snapshot.Bindings = append(snapshot.Bindings, old)
+				continue
+			}
+			return prior, rejected, fmt.Errorf("capacity cannot preserve established binding %s", binding.BindingUID)
+		}
+	}
+	sort.Slice(snapshot.Bindings, func(i, j int) bool { return snapshot.Bindings[i].BindingUID < snapshot.Bindings[j].BindingUID })
+	return snapshot, rejected, nil
+}
+
+func safeExpansion(old, candidate model.Binding) bool {
+	if old.Tombstone || candidate.Tombstone || old.ConfigurationPending || candidate.ConfigurationPending || old.Version != candidate.Version || old.ProjectUID != candidate.ProjectUID || old.ContextUID != candidate.ContextUID || old.ConsumerAddress != candidate.ConsumerAddress || old.ClusterAddress != candidate.ClusterAddress {
+		return false
+	}
+	for _, zone := range old.Zones {
+		found := false
+		for _, next := range candidate.Zones {
+			if zone.ZoneUID == next.ZoneUID && zone.Apex == next.Apex && !next.Pending {
+				found = true
+			}
+		}
+		if !found {
+			return false
+		}
+	}
+	return true
+}
+
+func (p *Planner) reportCapacity(ctx context.Context, bindings []dnsv1alpha1.DNSResolverBinding, rejected map[string]bool, now time.Time) error {
+	var failures []error
+	for i := range bindings {
+		binding := &bindings[i]
+		if binding.Spec.Placement.Region != p.Config.Region || binding.Spec.Placement.Shard != p.Config.Shard {
+			continue
+		}
+		base := binding.DeepCopy()
+		condition(&binding.Status.Conditions, "CapacityAvailable", !rejected[string(binding.UID)], "ShardCapacity", binding.Generation, now)
+		if !jsonEqualStatus(base.Status, binding.Status) {
+			if err := p.Client.Status().Patch(ctx, binding, client.MergeFromWithOptions(base, client.MergeFromWithOptimisticLock{})); err != nil {
+				failures = append(failures, err)
+			}
+		}
+	}
+	return errors.Join(failures...)
+}

internal/internaldns/platform/capacity_test.go

@@ -0,0 +1,76 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+
+package platform
+
+import (
+	"context"
+	"fmt"
+	"testing"
+	"time"
+
+	dnsv1alpha1 "go.miloapis.com/dns-operator/api/v1alpha1"
+	"go.miloapis.com/dns-operator/internal/internaldns/model"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"sigs.k8s.io/controller-runtime/pkg/client"
+)
+
+func capacityBinding(index int) model.Binding {
+	return model.Binding{BindingUID: fmt.Sprintf("binding-%d", index), ProjectUID: "project", ContextUID: fmt.Sprintf("context-%d", index), BindingGeneration: 1, ConfigurationRevision: 1, ConsumerAddress: fmt.Sprintf("fd53::%d", index), ClusterAddress: fmt.Sprintf("fd54::%d", index), Port: 53, Transports: []model.Transport{model.TransportUDP, model.TransportTCP}, Version: model.AccessVersion{Epoch: 1, Sequence: 1}}
+}
+
+func TestCapacityOverflowDoesNotBlockWithdrawal(t *testing.T) {
+	prior := model.ServingSnapshot{Region: "east", Shard: "a", ConfigurationEpoch: 1, ConfigurationRevision: 1, GeneratedAt: time.Now(), Bindings: []model.Binding{capacityBinding(1), capacityBinding(2)}}
+	candidate := prior
+	candidate.ConfigurationRevision++
+	candidate.Bindings = []model.Binding{capacityBinding(2), capacityBinding(3), capacityBinding(4)}
+	bounded, rejected, err := fitSnapshot(candidate, prior, 2, 512*1024)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(bounded.Bindings) != 2 || len(rejected) != 1 {
+		t.Fatalf("capacity result: %#v, rejected %#v", bounded.Bindings, rejected)
+	}
+	for _, binding := range bounded.Bindings {
+		if binding.BindingUID == "binding-1" {
+			t.Fatal("overflow prevented withdrawal")
+		}
+	}
+	if bounded.Bindings[0].BindingUID != "binding-2" {
+		t.Fatal("new allocations displaced an existing context")
+	}
+}
+
+func TestByteOverflowPreservesExistingContext(t *testing.T) {
+	prior := model.ServingSnapshot{Region: "east", Shard: "a", ConfigurationEpoch: 1, ConfigurationRevision: 1, GeneratedAt: time.Now(), Bindings: []model.Binding{capacityBinding(1)}}
+	candidate := prior
+	binding := capacityBinding(1)
+	for i := range 256 {
+		binding.Zones = append(binding.Zones, model.ZoneAttachment{ZoneUID: fmt.Sprintf("zone-%d", i), Apex: fmt.Sprintf("zone-%d.internal", i), RequiredPublicationEpoch: 1, RequiredPublicationRevision: 1})
+	}
+	binding.ConfigurationRevision++
+	candidate.Bindings = []model.Binding{binding, capacityBinding(2)}
+	bounded, rejected, err := fitSnapshot(candidate, prior, 2, 4096)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !rejected["binding-1"] || len(bounded.Bindings) != 2 || bounded.Bindings[0].ConfigurationRevision != 1 {
+		t.Fatalf("overflow interrupted an established context: %#v", bounded)
+	}
+}
+
+func TestCapacityReportDoesNotOverwriteAnotherShard(t *testing.T) {
+	ctx := context.Background()
+	binding := &dnsv1alpha1.DNSResolverBinding{ObjectMeta: metav1.ObjectMeta{Name: "binding", Namespace: "platform", UID: "binding-uid"}, Spec: dnsv1alpha1.DNSResolverBindingSpec{Placement: dnsv1alpha1.DNSResolverBindingPlacement{Region: "west", Shard: "b"}}, Status: dnsv1alpha1.DNSResolverBindingStatus{Conditions: []metav1.Condition{{Type: "CapacityAvailable", Status: metav1.ConditionFalse, Reason: "ShardCapacity", LastTransitionTime: metav1.Now()}}}}
+	cl := testClient(t, binding)
+	planner := Planner{Client: cl, Config: PlannerConfig{Region: "east", Shard: "a"}}
+	if err := planner.reportCapacity(ctx, []dnsv1alpha1.DNSResolverBinding{*binding}, nil, time.Now()); err != nil {
+		t.Fatal(err)
+	}
+	var updated dnsv1alpha1.DNSResolverBinding
+	if err := cl.Get(ctx, client.ObjectKeyFromObject(binding), &updated); err != nil {
+		t.Fatal(err)
+	}
+	if len(updated.Status.Conditions) != 1 || updated.Status.Conditions[0].Status != metav1.ConditionFalse {
+		t.Fatal("capacity from another shard was overwritten")
+	}
+}