Building an eBPF/XDP NAT-Based (Weighted) Least Connection Load Balancer from Scratch

// We could also include port information but we simplify
// and assume that both LB and Backend listen on the same port for requests
struct endpoint {
  __u32 ip;
};

struct backend {
  // Backend endpoint information (currently only IP, but could be extended with port or other metadata)
  struct endpoint endpoint;
  // Number of active connections to this backend, used for least-connections load balancing algorithm
  __u32 num_connections;
};

struct {
  __uint(type, BPF_MAP_TYPE_ARRAY);
  __uint(max_entries, NUM_BACKENDS);
  __type(key, __u32);
  __type(value, struct backend);
} backends SEC(".maps");

...
__u32 key = 0;
__u32 min_connections = (__u32) - 1; // Max value for unsigned int
struct backend *candidate_backend = NULL;
for (__u32 i = 0; i < NUM_BACKENDS; i++) {
  __u32 idx = i;
  struct backend *b = bpf_map_lookup_elem(&backends, &idx);
  if (b && b->num_connections < min_connections) {
    min_connections = b->num_connections;
    key = idx;
    candidate_backend = b;
  }
}
...

enum tcp_state {
  TCP_STATE_UNKNOWN = 0,          // Initial placeholder
  TCP_STATE_SYN_SEEN = 1,         // SYN seen from client, handshake not complete
  TCP_STATE_ESTABLISHED = 2,      // ACK seen from client, handshake complete
  TCP_STATE_FIN_FROM_CLIENT = 3,  // FIN seen from client
  TCP_STATE_FIN_FROM_BACKEND = 4, // FIN seen from backend
  TCP_STATE_FIN_BOTH = 5,         // FIN seen from both sides, waiting for final ACK
};

static __always_inline void update_tcp_conn_state(struct five_tuple_t five_tuple, 
                                                  struct connection *conn, 
                                                  struct tcphdr *tcp,
                                                  int direction /* 0 = client->backend, 1 = backend->client */) {
  switch (conn->state) {
  // Should be SYN from client to backend
  case TCP_STATE_UNKNOWN:
    if (direction == 0 && tcp->syn) {
      conn = update_state(&five_tuple, conn, TCP_STATE_SYN_SEEN);
    }
    return;

  // Should be ACK from client to backend completing the handshake
  case TCP_STATE_SYN_SEEN:
    if (direction == 0 && !tcp->syn && tcp->ack) {
      conn = update_state(&five_tuple, conn, TCP_STATE_ESTABLISHED);
      if (!conn) return;

      // Increment the connection count for the selected backend
      struct backend *b = bpf_map_lookup_elem(&backends, &conn->backend_index);
      if (b) {
        b->num_connections++;
        bpf_map_update_elem(&backends, &conn->backend_index, b, BPF_ANY);
        bpf_printk("Selected backend with IP %pI4 with current number of connections equal to %d", &b->endpoint.ip, b->num_connections);
      }
    }
    return;

  // From established connection we need to see FINs from both sides
  case TCP_STATE_ESTABLISHED:
  case TCP_STATE_FIN_FROM_CLIENT:
  case TCP_STATE_FIN_FROM_BACKEND:
    if (!tcp->fin) return;
    __u8 new_state;
    if (direction == 0) {
      new_state = (conn->state == TCP_STATE_FIN_FROM_BACKEND)
                    ? TCP_STATE_FIN_BOTH
                    : TCP_STATE_FIN_FROM_CLIENT;
    } else {
      new_state = (conn->state == TCP_STATE_FIN_FROM_CLIENT)
                    ? TCP_STATE_FIN_BOTH
                    : TCP_STATE_FIN_FROM_BACKEND;
    }

    conn = update_state(&five_tuple, conn, new_state);
    return;

  // After FIN from both sides, we wait for the final ACK or RST to clean up the connection
  case TCP_STATE_FIN_BOTH:
    if ((tcp->ack && !tcp->fin) || tcp->rst) {
      // Decrement the connection count for the selected backend
      struct backend *b = bpf_map_lookup_elem(&backends, &conn->backend_index);
      if (b) {
        if (b->num_connections > 0) b->num_connections--;
        bpf_map_update_elem(&backends, &conn->backend_index, b, BPF_ANY);
        bpf_printk("Connection closed, backend %pI4 now has %d connections", &b->endpoint.ip, b->num_connections);
      }
      bpf_map_delete_elem(&statetrack, &five_tuple);
    }
    return;

  default:
    return;
  }
}

SEC("xdp")
int xdp_load_balancer(struct xdp_md *ctx) {
  ...
  struct endpoint *out = bpf_map_lookup_elem(&conntrack, &in);
  if (!out) {
    //bpf_printk("Packet from client..");
    ...
    // Update state for both existing and new connections on packets from clients
    update_tcp_conn_state(five_tuple, conn_ptr, tcp, 0);
    ...
  } else {
    //bpf_printk("Packet from backend..");
    ...
    // Update state on packet from backend
    update_tcp_conn_state(out_loadbalancer, conn, tcp, 1);
    ...
  }
  ...
}

struct endpoint {
  __u32 ip;
};

struct backend {
  // Backend endpoint information (currently only IP, but could be extended with port or other metadata)
  struct endpoint endpoint;
  // Number of active connections to this backend, used for least-connections load balancing algorithm
  __u32 num_connections;
  // Backend weight for weighted load balancing algorithms
  __u32 weight;
};

...
__u32 key = 0;
__u32 min_score = (__u32) - 1; // Max value for unsigned int
struct backend *candidate_backend = NULL;
for (__u32 i = 0; i < NUM_BACKENDS; i++) {
  __u32 idx = i;
  struct backend *b = bpf_map_lookup_elem(&backends, &idx);
  if (b && b->weight > 0) {
    // Calculate score: (connections * scale) / weight
    // Higher weight reduces the score, making the backend more likely to be picked.
    // We add 1 to num_connections so that weight matters even at 0 connections.
    __u32 score = ((b->num_connections + 1) * 1024) / b->weight;

    if (score < min_score) {
      min_score = score;
      key = idx;
      candidate_backend = b;
    }
  }
}
...