Building an eBPF/XDP NAT-Based (Weighted) Round Robin Load Balancer from Scratch

struct {
  __uint(type, BPF_MAP_TYPE_ARRAY);
  __uint(max_entries, 1);
  __type(key, __u32);
  __type(value, __u32); // Holds the next backend ID to sent traffic to
} scheduler_state SEC(".maps");

SEC("xdp")
int xdp_load_balancer(struct xdp_md *ctx) {
  ...
  // Lookup conntrack (connection tracking) information:
  // - No Connection:    client request
  // - Connection exist: backend response
  struct endpoint *out = bpf_map_lookup_elem(&conntrack, &in);
  if (!out) {
    // bpf_printk("Packet from client..");
    // Check for an existing connection
    ...
    __u32 *backend_idx = bpf_map_lookup_elem(&statetrack, &five_tuple);
    if (backend_idx) {
      // Existing connection found in statetrack map - proceed with the same backend
      backend = bpf_map_lookup_elem(&backends, backend_idx);
      ...
    } else {
      // New connection - select backend with round-robin scheduling
      __u32 zero = 0;
      __u32 *curr_idx = bpf_map_lookup_elem(&scheduler_state, &zero);
      if (!curr_idx) {
        return XDP_ABORTED;
      }
      __u32 key = *curr_idx;
      backend = bpf_map_lookup_elem(&backends, &key);
      if (!backend) {
        return XDP_ABORTED;
      }
      bpf_printk("Selected backend index: %d with IP: %pI4", key, &backend->ip);

      // Increment the count to point to the next backend and update the entry in scheduler_state map
      // Perform modulo operation to get a number in the range of backend indexes
      __u32 next_idx = (key + 1) % NUM_BACKENDS;
      if (bpf_map_update_elem(&scheduler_state, &zero, &next_idx, BPF_ANY) != 0) {
        return XDP_ABORTED;
      }
      ...
    }
  } else {
    // bpf_printk("Packet from backend..");
    // Packets from backends are just sent back to client - nothing else to do
    ...
  }
  ...
}

struct backend {
  // Backend endpoint information (currently only IP, but could be extended with port or other metadata)
  struct endpoint endpoint;
  // Backend weight for weighted load balancing algorithms
  __u32 weight;
  // Number of connections handled by this backend
  __u32 used_count; 
};

SEC("xdp")
int xdp_load_balancer(struct xdp_md *ctx) {
  ...
  // Lookup conntrack (connection tracking) information:
  // - No Connection:    client request
  // - Connection exist: backend response
  struct endpoint *out = bpf_map_lookup_elem(&conntrack, &in);
  if (!out) {
    //bpf_printk("Packet from client..");
    // Check for existing connections
    ...
    __u32 *backend_idx = bpf_map_lookup_elem(&statetrack, &five_tuple);
    if (backend_idx) {
      // Existing connection found in statetrack map - update state and proceed with the same backend
      backend = bpf_map_lookup_elem(&backends, backend_idx);
      ...
    } else {
      // No existing connection found in statetrack map
      // Select a backend using weighted round robin algorithm
      __u32 zero = 0;
      __u32 *curr_idx = bpf_map_lookup_elem(&scheduler_state, &zero);
      if (!curr_idx) {
        return XDP_ABORTED;
      }
      __u32 selected = *curr_idx;
      backend = bpf_map_lookup_elem(&backends, &selected);
      if (!backend) {
        return XDP_ABORTED;
      }

      // Increment used_count as we are adding a new connection to this backend
      backend->used_count += 1;

      // Check whether used_count is equal to the backend's weight
      // This means that this backend has handled the number of connections equal to its weight and 
      // we can move to the next backend for the future connections
      if (backend->used_count >= backend->weight) { 
        backend->used_count = 0;                         // Rest used_count to 0 when it's equal its weight
        __u32 next_idx = (selected + 1) % NUM_BACKENDS;  // Increment the index to point to the next backend
        if (bpf_map_update_elem(&scheduler_state, &zero, &next_idx, BPF_ANY) != 0) {
          return XDP_ABORTED;
        }
      }
      bpf_printk("Selected backend index: %d with IP: %pI4", selected, &backend->endpoint.ip);
      ...
    }
    ...
  } else {
    //bpf_printk("Packet from backend..");
    // Packets from backends are just sent back to client - nothing else to do
    ...
  }
  ...
}