Add location-aware routing (#221)

* Add location-aware routing for proxy targets and provisioning

Ling proxy now prefers targets on kings matching its preferred location,
falling back to all targets when none match. Council provisioner prefers
kings matching a service's preferred location when assigning ports.

Closes #14

* Fix lint issues: golines, wsl, unparam

Author: firecow

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
 }