feat: Order Caddy upstreams and dns nearest results by Corrosion median RTT with machine

internal/machine/caddyconfig/caddyfile.go

@@ -14,6 +14,7 @@ import (
 	"text/template"
 	"time"
 
+	"github.com/psviderski/uncloud/internal/machine/rtt"
 	"github.com/psviderski/uncloud/internal/machine/store"
 	"github.com/psviderski/uncloud/pkg/api"
 )
@@ -70,8 +71,10 @@ type CaddyfileGenerator struct {
 	machineID string
 	// machineName is the human-friendly name of the machine.
 	machineName string
-	validator   CaddyfileValidator
-	log         *slog.Logger
+	// rttCache provides RTT lookup by machine ID. Nil-safe: RTTFor returns UnknownRTT on nil.
+	rttCache  *rtt.Cache
+	validator CaddyfileValidator
+	log       *slog.Logger
 }
 
 // CaddyfileValidator is an interface for validating Caddyfile configurations.
@@ -80,14 +83,16 @@ type CaddyfileValidator interface {
 }
 
 func NewCaddyfileGenerator(
-	machineID, machineName string, validator CaddyfileValidator, log *slog.Logger,
+	machineID, machineName string, rttCache *rtt.Cache,
+	validator CaddyfileValidator, log *slog.Logger,
 ) *CaddyfileGenerator {
 	if log == nil {
 		log = slog.Default()
 	}
 	return &CaddyfileGenerator{
 		machineID:   machineID,
 		machineName: machineName,
+		rttCache:    rttCache,
 		validator:   validator,
 		log:         log,
 	}
@@ -119,7 +124,7 @@ func (g *CaddyfileGenerator) Generate(
 	// The service name and creation time tiebreakers keep the generated Caddyfile stable across regenerations.
 	slices.SortStableFunc(records, func(a, b store.ContainerRecord) int {
 		return cmp.Or(
-			g.localMachineRank(a.MachineID)-g.localMachineRank(b.MachineID),
+			cmp.Compare(g.rttCache.RTTFor(a.MachineID), g.rttCache.RTTFor(b.MachineID)),
 			strings.Compare(a.Container.ServiceName(), b.Container.ServiceName()),
 			a.Container.CreatedTime().Compare(b.Container.CreatedTime()),
 		)
@@ -249,15 +254,6 @@ func (g *CaddyfileGenerator) Generate(
 	return caddyfileHeader + "\n" + caddyfile, nil
 }
 
-// localMachineRank returns 0 if the given machineID matches the local machine and 1 otherwise.
-// Useful for sorting containers running locally first.
-func (g *CaddyfileGenerator) localMachineRank(machineID string) int {
-	if g.machineID == machineID {
-		return 0
-	}
-	return 1
-}
-
 func (g *CaddyfileGenerator) generateBaseFromPorts(containers []api.ServiceContainer) (string, error) {
 	httpHostUpstreams, httpsHostUpstreams := httpUpstreamsFromPorts(containers)
 

internal/machine/caddyconfig/caddyfile_test.go

@@ -11,6 +11,7 @@ import (
 	"github.com/docker/docker/api/types/container"
 	"github.com/docker/docker/api/types/network"
 	"github.com/psviderski/uncloud/internal/machine/docker"
+	"github.com/psviderski/uncloud/internal/machine/rtt"
 	"github.com/psviderski/uncloud/internal/machine/store"
 	"github.com/psviderski/uncloud/pkg/api"
 	"github.com/stretchr/testify/assert"
@@ -209,7 +210,11 @@ http://app.example.com {
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			// Validator is not expected to be called in these tests.
-			generator := NewCaddyfileGenerator("test-machine-id", "test-machine", nil, nil)
+			// Provide a cache with the local machine so RTT-based sorting works.
+			rttCache := rtt.NewCacheWithStats("test-machine-id", map[string]rtt.Stats{
+				"test-machine-id": {Median: 0},
+			})
+			generator := NewCaddyfileGenerator("test-machine-id", "test-machine", rttCache, nil, nil)
 
 			config, err := generator.Generate(ctx, tt.containers, true)
 
@@ -853,7 +858,7 @@ valid.example.com {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			generator := NewCaddyfileGenerator("test-machine-id", "test-machine", validator, nil)
+			generator := NewCaddyfileGenerator("test-machine-id", "test-machine", nil, validator, nil)
 
 			config, err := generator.Generate(ctx, tt.containers, true)
 
@@ -995,7 +1000,7 @@ http://api.example.com {
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			// Validator is not expected to be called in these tests.
-			generator := NewCaddyfileGenerator("test-machine-id", "test-machine", nil, nil)
+			generator := NewCaddyfileGenerator("test-machine-id", "test-machine", nil, nil, nil)
 
 			config, err := generator.Generate(ctx, tt.containers, false)
 			require.NoError(t, err)

internal/machine/caddyconfig/controller.go

@@ -12,6 +12,7 @@ import (
 	"strings"
 
 	"github.com/psviderski/uncloud/internal/fs"
+	"github.com/psviderski/uncloud/internal/machine/rtt"
 	"github.com/psviderski/uncloud/internal/machine/store"
 	"github.com/psviderski/uncloud/pkg/api"
 )
@@ -31,6 +32,7 @@ type Controller struct {
 	generator     *CaddyfileGenerator
 	client        *CaddyAdminClient
 	store         *store.Store
+	rttCache      *rtt.Cache
 	log           *slog.Logger
 	// lastFingerprint caches the fingerprint of the containers used to generate the latest successfully loaded
 	// Caddyfile. nil means it hasn't been loaded yet or the last load failed.
@@ -56,7 +58,7 @@ func (f containerFingerprint) Equal(other containerFingerprint) bool {
 		f.CaddyConfig == other.CaddyConfig
 }
 
-func NewController(machineID, configDir, adminSock string, store *store.Store) (*Controller, error) {
+func NewController(machineID, configDir, adminSock string, store *store.Store, rttCache *rtt.Cache) (*Controller, error) {
 	if err := os.MkdirAll(configDir, 0o750); err != nil {
 		return nil, fmt.Errorf("create directory for Caddy configuration '%s': %w", configDir, err)
 	}
@@ -73,6 +75,7 @@ func NewController(machineID, configDir, adminSock string, store *store.Store) (
 		caddyfilePath: filepath.Join(configDir, "Caddyfile"),
 		client:        client,
 		store:         store,
+		rttCache:      rttCache,
 		log:           log,
 	}, nil
 }
@@ -87,7 +90,7 @@ func (c *Controller) Run(ctx context.Context) error {
 	} else {
 		machineName = m.Name
 	}
-	c.generator = NewCaddyfileGenerator(c.machineID, machineName, c.client, c.log)
+	c.generator = NewCaddyfileGenerator(c.machineID, machineName, c.rttCache, c.client, c.log)
 
 	containers, changes, err := c.store.SubscribeContainers(ctx)
 	if err != nil {

internal/machine/dns/resolver.go

@@ -5,7 +5,6 @@ import (
 	"fmt"
 	"log/slog"
 	"maps"
-	"net/netip"
 	"slices"
 	"sync"
 	"time"
@@ -17,8 +16,8 @@ import (
 // to their IP addresses.
 type ClusterResolver struct {
 	store *store.Store
-	// serviceIPs maps service names to container IPs.
-	serviceIPs map[string][]netip.Addr
+	// serviceIPs maps service names to resolved container IPs with machine metadata.
+	serviceIPs map[string][]ResolvedIP
 	// mu protects the serviceIPs map.
 	mu sync.RWMutex
 	// lastUpdate tracks when records were last updated.
@@ -30,7 +29,7 @@ type ClusterResolver struct {
 func NewClusterResolver(store *store.Store) *ClusterResolver {
 	return &ClusterResolver{
 		store:      store,
-		serviceIPs: make(map[string][]netip.Addr),
+		serviceIPs: make(map[string][]ResolvedIP),
 		log:        slog.With("component", "dns-resolver"),
 	}
 }
@@ -70,7 +69,7 @@ func (r *ClusterResolver) Run(ctx context.Context) error {
 
 // updateServiceIPs processes container records and updates the serviceIPs map.
 func (r *ClusterResolver) updateServiceIPs(containers []store.ContainerRecord) {
-	newServiceIPs := make(map[string][]netip.Addr, len(r.serviceIPs))
+	newServiceIPs := make(map[string][]ResolvedIP, len(r.serviceIPs))
 
 	containersCount := 0
 	for _, record := range containers {
@@ -93,23 +92,24 @@ func (r *ClusterResolver) updateServiceIPs(containers []store.ContainerRecord) {
 			continue
 		}
 
-		newServiceIPs[ctr.ServiceName()] = append(newServiceIPs[ctr.ServiceName()], ip)
+		resolved := ResolvedIP{Addr: ip, MachineID: record.MachineID}
+		newServiceIPs[ctr.ServiceName()] = append(newServiceIPs[ctr.ServiceName()], resolved)
 		// Also add the service ID as a valid lookup.
-		newServiceIPs[ctr.ServiceID()] = append(newServiceIPs[ctr.ServiceID()], ip)
+		newServiceIPs[ctr.ServiceID()] = append(newServiceIPs[ctr.ServiceID()], resolved)
 
 		// Add <machine-id>.m.<service-name> as a lookup
 		serviceNameWithMachineID := record.MachineID + ".m." + ctr.ServiceName()
-		newServiceIPs[serviceNameWithMachineID] = append(newServiceIPs[serviceNameWithMachineID], ip)
+		newServiceIPs[serviceNameWithMachineID] = append(newServiceIPs[serviceNameWithMachineID], resolved)
 
 		containersCount++
 	}
 
-	// Sort each service's IPs so they have a deterministic order for comparison.
+	// Sort each service's resolved IPs by address for deterministic order and comparison.
 	for _, ips := range newServiceIPs {
-		slices.SortFunc(ips, func(a, b netip.Addr) int { return a.Compare(b) })
+		slices.SortFunc(ips, func(a, b ResolvedIP) int { return a.Addr.Compare(b.Addr) })
 	}
 	// Skip the swap when the services or their container IPs haven't changed.
-	if maps.EqualFunc(r.serviceIPs, newServiceIPs, slices.Equal[[]netip.Addr]) {
+	if maps.EqualFunc(r.serviceIPs, newServiceIPs, slices.Equal[[]ResolvedIP]) {
 		return
 	}
 
@@ -121,8 +121,8 @@ func (r *ClusterResolver) updateServiceIPs(containers []store.ContainerRecord) {
 	r.log.Info("DNS records updated.", "services", len(newServiceIPs)/3, "containers", containersCount)
 }
 
-// Resolve returns IP addresses of the service containers.
-func (r *ClusterResolver) Resolve(serviceName string) []netip.Addr {
+// Resolve returns resolved IPs of the service containers.
+func (r *ClusterResolver) Resolve(serviceName string) []ResolvedIP {
 	r.mu.RLock()
 	defer r.mu.RUnlock()
 
@@ -132,7 +132,7 @@ func (r *ClusterResolver) Resolve(serviceName string) []netip.Addr {
 	}
 
 	// Return a copy of the IPs slice to prevent modification of the original.
-	ipsCopy := make([]netip.Addr, len(ips))
+	ipsCopy := make([]ResolvedIP, len(ips))
 	copy(ipsCopy, ips)
 
 	return ipsCopy

internal/machine/dns/server.go

@@ -1,6 +1,7 @@
 package dns
 
 import (
+	"cmp"
 	"context"
 	"errors"
 	"fmt"
@@ -15,6 +16,7 @@ import (
 	"time"
 
 	"github.com/miekg/dns"
+	"github.com/psviderski/uncloud/internal/machine/rtt"
 	"github.com/psviderski/uncloud/internal/metrics"
 )
 
@@ -31,19 +33,25 @@ const (
 	forwardingTimeout = 3 * time.Second
 )
 
-// Resolver is an interface for resolving service names to IP addresses.
+// ResolvedIP pairs an IP address with the machine ID where the container is running.
+type ResolvedIP struct {
+	Addr      netip.Addr
+	MachineID string
+}
+
+// Resolver is an interface for resolving service names to IP addresses with machine metadata.
 type Resolver interface {
-	// Resolve returns a list of IP addresses of the service containers.
+	// Resolve returns a list of resolved IPs of the service containers.
 	// An empty list is returned if no service is found.
-	Resolve(serviceName string) []netip.Addr
+	Resolve(serviceName string) []ResolvedIP
 }
 
 // Server is an embedded internal DNS server for service discovery and forwarding external queries
 // to upstream DNS servers.
 type Server struct {
 	listenAddr      netip.Addr
-	localSubnet     netip.Prefix
 	resolver        Resolver
+	rttCache        *rtt.Cache
 	upstreamServers []netip.AddrPort
 
 	udpServer        *dns.Server
@@ -56,7 +64,7 @@ type Server struct {
 // NewServer creates a new DNS server with the given configuration.
 // If upstreams is nil, nameservers from /etc/resolv.conf will be used. An empty upstreams list means to only resolve
 // internal DNS queries and not forward any external queries.
-func NewServer(listenAddr netip.Addr, localSubnet netip.Prefix, resolver Resolver, upstreams []netip.AddrPort) (*Server, error) {
+func NewServer(listenAddr netip.Addr, resolver Resolver, upstreams []netip.AddrPort, rttCache *rtt.Cache) (*Server, error) {
 	if !listenAddr.IsValid() {
 		return nil, fmt.Errorf("invalid listen address: %s", listenAddr)
 	}
@@ -90,8 +98,8 @@ func NewServer(listenAddr netip.Addr, localSubnet netip.Prefix, resolver Resolve
 
 	return &Server{
 		listenAddr:       listenAddr,
-		localSubnet:      localSubnet,
 		resolver:         resolver,
+		rttCache:         rttCache,
 		upstreamServers:  upstreams,
 		forwardSemaphore: make(chan struct{}, maxConcurrentForwards),
 		log:              slog.With("component", "dns-server"),
@@ -312,23 +320,22 @@ func (s *Server) handleAQuery(name string) []dns.RR {
 		// and nothing additional to do for round-robin (mode == "rr").
 
 		if mode == "nearest" {
-			// Sort IPs on local subnet to the top.
-			slices.SortFunc(ips, func(a, b netip.Addr) int {
-				aIsLocal := s.localSubnet.Contains(a)
-				bIsLocal := s.localSubnet.Contains(b)
-				if aIsLocal && !bIsLocal {
-					return -1
-				} else if bIsLocal && !aIsLocal {
-					return 1
-				}
-				return 0
+			// Sort by RTT using proximity data. Local machine containers get RTT 0.
+			slices.SortStableFunc(ips, func(a, b ResolvedIP) int {
+				return cmp.Compare(s.rttCache.RTTFor(a.MachineID), s.rttCache.RTTFor(b.MachineID))
 			})
 		}
 	}
 
+	// Extract addresses from resolved IPs for DNS response.
+	addrs := make([]netip.Addr, len(ips))
+	for i, r := range ips {
+		addrs[i] = r.Addr
+	}
+
 	// Create A records for each IP.
-	records := make([]dns.RR, 0, len(ips))
-	for _, ip := range ips {
+	records := make([]dns.RR, 0, len(addrs))
+	for _, ip := range addrs {
 		records = append(records, &dns.A{
 			Hdr: dns.RR_Header{
 				Name:   name,