# The Load Balancer

> Distribute incoming requests evenly across available servers.

Canonical URL: <https://datadriven.io/problems/the_load_balancer>

Domain: Python · Difficulty: medium · Seniority: L4

## Problem

Implement a LoadBalancer via an ops list. Operations: ['add_server', id], ['remove_server', id], ['get_server']. get_server returns the next server in round-robin over the current servers in the order they were added. Return the result list: None for add/remove, the server id for get_server. If get_server is called with zero servers, the behavior is undefined (test data will not hit this).

## Worked solution and explanation

### Why this problem exists in real interviews

This tests **class design, state management, and modular arithmetic**. Round-robin load balancing is a real systems concept, and interviewers use it to check whether candidates can manage mutable state across method calls while handling edge cases like empty server lists.

---

### Break down the requirements

#### Step 1: Maintain a list of active servers

Store server IDs in a list. Additions append, removals filter out the ID.

#### Step 2: Track the current rotation index

A counter tracks which server to return next. It increments with each `get_server` call.

#### Step 3: Wrap around using modular arithmetic

Use `index % len(servers)` to cycle back to the beginning after reaching the end of the list.

#### Step 4: Raise on empty

If no servers are registered, `get_server` must raise an exception rather than returning None or crashing on division by zero.

---

### The solution

**Round-robin rotation with modular index**

```python
class LoadBalancer:
    def __init__(self):
        self.servers = []
        self.index = 0
    def add_server(self, server_id):
        self.servers.append(server_id)
    def remove_server(self, server_id):
        self.servers.remove(server_id)
        if self.index >= len(self.servers) and len(self.servers) > 0:
            self.index = self.index % len(self.servers)
    def get_server(self):
        if not self.servers:
            raise Exception("No servers available")
        server = self.servers[self.index % len(self.servers)]
        self.index = (self.index + 1) % len(self.servers)
        return server
```

> **Time and Space Complexity**
>
> **Time:** O(1) for `add_server` and `get_server`. O(n) for `remove_server` due to list removal.
> 
> **Space:** O(n) where n is the number of registered servers.

> **Interviewers Watch For**
>
> How do you handle index drift after removal? If the current index points past the end of the shortened list, you must reset it. Missing this causes IndexError in production.

> **Common Pitfall**
>
> Using `self.index += 1` without modular wrapping. The index grows unbounded, which technically works with `%` at read time but is a code smell that signals the candidate is not thinking about state invariants.

---

## Common follow-up questions

- How would you implement weighted round-robin? _(Tests extending the rotation to account for server capacity differences.)_
- What if add and remove are called concurrently from multiple threads? _(Tests thread-safety awareness: locks, atomic operations, or lock-free data structures.)_
- How would you add health checks that auto-remove unresponsive servers? _(Tests integrating monitoring with state management.)_
- Why might round-robin be a poor choice for long-lived connections? _(Tests understanding of connection-aware balancing strategies like least-connections.)_

## Related

- [All practice problems](https://datadriven.io/problems)
- [Mock interview mode](https://datadriven.io/interview/the_load_balancer)
- [Python Interview Questions](https://datadriven.io/python-interview-questions)
- [Data Engineering Interview Prep Guide](https://datadriven.io/data-engineer-interview-prep)
- [Daily Challenge](https://datadriven.io/daily)

---

Source: DataDriven (https://datadriven.io). 100% free data engineering interview prep. Live code execution against Postgres 16, Python 3.11, and Spark sandboxes. No paywall, no premium tier, no signup gate.