Add location-aware routing

internal/council/provisioner.go

@@ -82,7 +82,8 @@ func collectUsedPorts(king *state.King, currentState *state.State) map[int]bool
 	return used
 }
 
-// ProvisionService assigns the first available king port to a service.
+// ProvisionService assigns the first available king port to a service,
+// preferring kings that match the service's PreferredLocation.
 func ProvisionService(currentState *state.State, service *state.Service) {
 	available := AvailableKingPorts(currentState)
 	if len(available) == 0 {
@@ -91,19 +92,30 @@ func ProvisionService(currentState *state.State, service *state.Service) {
 		return
 	}
 
-	first := available[0]
-	remotePort := first.Ports[0]
+	selected := available[0]
 
-	service.Host = &first.King.Host
+	if service.PreferredLocation != "" {
+		for _, candidate := range available {
+			if candidate.King.Location == service.PreferredLocation {
+				selected = candidate
+
+				break
+			}
+		}
+	}
+
+	remotePort := selected.Ports[0]
+
+	service.Host = &selected.King.Host
 	service.RemotePort = &remotePort
-	service.BindPort = &first.King.BindPort
+	service.BindPort = &selected.King.BindPort
 
 	currentState.Revision++
 
 	slog.Info("Provisioned service",
 		"name", service.Name,
-		"host", first.King.Host,
-		"bind_port", first.King.BindPort,
+		"host", selected.King.Host,
+		"bind_port", selected.King.BindPort,
 		"remote_port", remotePort,
 	)
 }

internal/council/provisioner_test.go

@@ -568,3 +568,101 @@ func TestProvision_NoReprovisioning_WhenNoAvailableKing(t *testing.T) {
 		t.Fatalf("expected KingReady to be false")
 	}
 }
+
+func newTestKingWithLocation(
+	host, ports string, bindPort int, location string,
+) state.King {
+	return state.King{
+		BindPort:     bindPort,
+		Host:         host,
+		Ports:        ports,
+		ShuttingDown: false,
+		Beat:         0,
+		Location:     location,
+		CertPEM:      "",
+	}
+}
+
+func TestProvisionService_PrefersMatchingLocation(t *testing.T) {
+	t.Parallel()
+
+	currentState := &state.State{
+		Revision: 0,
+		Kings: []state.King{
+			newTestKingWithLocation(testHostA, "5000-5001", 2333, "us-east"),
+			newTestKingWithLocation(testHostB, "6000-6001", 2334, "eu-west"),
+		},
+		Services: []state.Service{
+			{
+				ServiceID:         "svc-1",
+				Name:              "alpha",
+				Token:             "",
+				LingID:            "",
+				PreferredLocation: "eu-west",
+				LingReady:         false,
+				KingReady:         false,
+				Host:              nil,
+				BindPort:          nil,
+				RemotePort:        nil,
+			},
+		},
+		Lings: nil,
+	}
+
+	council.ProvisionService(currentState, &currentState.Services[0])
+
+	svc := currentState.Services[0]
+	assertServiceHost(t, svc, testHostB)
+	assertServiceBindPort(t, svc, 2334)
+	assertServiceRemotePort(t, svc, 6000)
+}
+
+func TestProvisionService_FallsBackWhenPreferredLocationFull(t *testing.T) {
+	t.Parallel()
+
+	hostB := testHostB
+	bindPortB := 2334
+	port6000 := 6000
+
+	currentState := &state.State{
+		Revision: 0,
+		Kings: []state.King{
+			newTestKingWithLocation(testHostA, "5000-5001", 2333, "us-east"),
+			newTestKingWithLocation(testHostB, "6000-6000", 2334, "eu-west"),
+		},
+		Services: []state.Service{
+			{
+				ServiceID:         "existing",
+				Name:              "existing",
+				Token:             "",
+				LingID:            "",
+				PreferredLocation: "",
+				LingReady:         false,
+				KingReady:         false,
+				Host:              &hostB,
+				BindPort:          &bindPortB,
+				RemotePort:        &port6000,
+			},
+			{
+				ServiceID:         "svc-1",
+				Name:              "alpha",
+				Token:             "",
+				LingID:            "",
+				PreferredLocation: "eu-west",
+				LingReady:         false,
+				KingReady:         false,
+				Host:              nil,
+				BindPort:          nil,
+				RemotePort:        nil,
+			},
+		},
+		Lings: nil,
+	}
+
+	council.ProvisionService(currentState, &currentState.Services[1])
+
+	svc := currentState.Services[1]
+	assertServiceHost(t, svc, testHostA)
+	assertServiceBindPort(t, svc, 2333)
+	assertServiceRemotePort(t, svc, 5000)
+}

internal/ling/ling.go

@@ -253,7 +253,7 @@ func run(
 			stateMutex.Unlock()
 
 			OnStateChanged(
-				ctx, stateSnapshot, lingID,
+				ctx, stateSnapshot, lingID, preferredLocation,
 				tunnelMap, tunnelCli, tcpProxies,
 			)
 		},
@@ -435,13 +435,14 @@ func OnStateChanged(
 	ctx context.Context,
 	stateSnapshot *state.State,
 	lingID string,
+	preferredLocation string,
 	tunnelMap map[string]string,
 	tunnelCli *TunnelClient,
 	tcpProxies map[string]*TCPProxy,
 ) {
 	kings := buildKingIndex(stateSnapshot)
 	updateTunnelConnections(ctx, stateSnapshot, lingID, tunnelMap, tunnelCli, kings)
-	updateProxyTargets(stateSnapshot, tcpProxies, kings)
+	updateProxyTargets(stateSnapshot, preferredLocation, tcpProxies, kings)
 }
 
 // KingGroup holds the connection details for a group of services on a single king.
@@ -454,16 +455,21 @@ type KingGroup struct {
 
 // KingIndex stores health and certificate info for a king address.
 type KingIndex struct {
-	healthy bool
-	certPEM string
+	healthy  bool
+	certPEM  string
+	location string
 }
 
 func buildKingIndex(stateSnapshot *state.State) map[string]KingIndex {
 	index := make(map[string]KingIndex, len(stateSnapshot.Kings))
 
 	for _, king := range stateSnapshot.Kings {
 		addr := net.JoinHostPort(king.Host, strconv.Itoa(king.BindPort))
-		index[addr] = KingIndex{healthy: !king.ShuttingDown, certPEM: king.CertPEM}
+		index[addr] = KingIndex{
+			healthy:  !king.ShuttingDown,
+			certPEM:  king.CertPEM,
+			location: king.Location,
+		}
 	}
 
 	return index
@@ -551,12 +557,13 @@ func updateTunnelConnections(
 
 func updateProxyTargets(
 	stateSnapshot *state.State,
+	preferredLocation string,
 	tcpProxies map[string]*TCPProxy,
 	kings map[string]KingIndex,
 ) {
 	for proxyName, proxy := range tcpProxies {
 		targets := collectProxyTargets(
-			stateSnapshot, proxyName, kings,
+			stateSnapshot, proxyName, preferredLocation, kings,
 		)
 		proxy.updateTargets(targets)
 	}
@@ -589,20 +596,36 @@ func isServiceEligibleForProxy(
 func collectProxyTargets(
 	stateSnapshot *state.State,
 	proxyName string,
+	preferredLocation string,
 	kings map[string]KingIndex,
 ) []ProxyTarget {
-	targets := make([]ProxyTarget, 0, len(stateSnapshot.Services))
+	var allTargets []ProxyTarget
+
+	var preferredTargets []ProxyTarget
 
 	for _, svc := range stateSnapshot.Services {
 		if !isServiceEligibleForProxy(svc, proxyName, kings) {
 			continue
 		}
 
-		targets = append(targets, ProxyTarget{
+		target := ProxyTarget{
 			host:       *svc.Host,
 			remotePort: *svc.RemotePort,
-		})
+		}
+
+		allTargets = append(allTargets, target)
+
+		if preferredLocation != "" {
+			kingAddr := net.JoinHostPort(*svc.Host, strconv.Itoa(*svc.BindPort))
+			if king, exists := kings[kingAddr]; exists && king.location == preferredLocation {
+				preferredTargets = append(preferredTargets, target)
+			}
+		}
+	}
+
+	if len(preferredTargets) > 0 {
+		return preferredTargets
 	}
 
-	return targets
+	return allTargets
 }