Implement a Load-Balanced Session Manager with Session Migration
Company: Palantir
Role: Software Engineer
Category: Software Engineering Fundamentals
Difficulty: medium
Interview Round: Technical Screen
Implement a session manager that assigns active sessions to a fixed set of servers. Each session adds a load of 1 to the server it is on, so a server's load is the number of active sessions assigned to it. After every public operation, the highest and the lowest server loads may differ by at most 1, and the manager may move existing sessions between servers to keep it that way. The core task is the allocation itself (Part 1); ending sessions with rebalancing (Part 2) and the update-time target (Part 3) extend it.
**Interface**
```python
class SessionManager:
def __init__(self, servers: list[str]): ...
def start_session(self, session_id: str) -> None: ...
def end_session(self, session_id: str) -> None: ...
def get_allocation(self) -> dict[str, set[str]]: ...
```
**Balancing rule.** After every public operation, `max(load) - min(load) <= 1` over all servers. Loads `[3, 3, 2]`, `[2, 2, 2]` and `[1, 0, 0]` are valid; `[3, 3, 1]` is not.
**Example 1: starting sessions and rebalancing**
```python
manager = SessionManager(["s1", "s2", "s3"])
for s in ["a", "b", "c", "d", "e"]:
manager.start_session(s)
# One valid allocation, loads [2, 2, 1]:
# {"s1": {"a", "d"}, "s2": {"b", "e"}, "s3": {"c"}}
manager.start_session("a") # "a" is already active: the allocation is unchanged
manager.end_session("c") # from the allocation above; removing "c" alone leaves loads [2, 2, 0]
# One valid final allocation, loads [1, 2, 1]:
# {"s1": {"a"}, "s2": {"b", "e"}, "s3": {"d"}}
```
Other allocations that satisfy every requirement are equally valid.
**Example 2: missing sessions and reused IDs**
```python
manager = SessionManager(["s1"])
manager.end_session("x") # allocation: {"s1": set()}
manager.start_session("x") # allocation: {"s1": {"x"}}
manager.end_session("x") # allocation: {"s1": set()}
manager.start_session("x") # allocation: {"s1": {"x"}}
```
### Constraints and Clarifications
- `servers` is nonempty and contains unique server IDs. The set of servers never changes, and every server starts with no sessions.
- `start_session` on a session that is already active does nothing. Otherwise it assigns the new session to a server; when several allocations are valid, any of them is acceptable.
- `end_session` on a session that is not active does nothing. Otherwise it removes the session and moves other sessions if necessary to restore the rule. A move never creates or deletes a session. An ended session ID may be started again later.
- `get_allocation` returns every server, including servers with no sessions, and lists every active session exactly once. It returns an independent snapshot: changing the returned dictionary or any of its sets must not change the manager's state.
- Let S be the number of servers and A the number of active sessions. `start_session` and `end_session` must not scan all servers or all active sessions on every call; the target is O(log S) time or better per update. The constructor and `get_allocation` are exempt, and a snapshot may take O(S + A).
- Because valid allocations are not unique, tests check properties rather than one exact allocation: the load difference is at most 1, the set of active session IDs is correct, each active session appears exactly once, and every configured server is present.
- Scenarios to test: no sessions started; several distinct sessions started; an active ID started again; an unknown session ended; an end that needs no migration; an end that needs a migration; ending a session that was moved earlier; ending every session; modifying a returned snapshot.
### Clarifying Questions
- Is any rebalancing that restores the rule acceptable, or should the manager move as few sessions as possible?
- When a session has to move, may it be any session on the source server, or should some sessions (for example, newer ones) be preferred?
- Should the constructor reject an empty or duplicated server list, or may it rely on the stated guarantees?
- Is the manager called from a single thread?
### Part 1 — Start sessions and report the allocation
Implement the constructor, `start_session` and `get_allocation` so that the balancing rule holds after every call. Walk through Example 1 up to the repeated start of `"a"`.
```hint Where a new session can go
Before a start, the loads already differ by at most one. Work out which servers can take one more session without breaking the rule, and whether a start ever has to move an existing session.
```
#### What This Part Should Cover
- Which server receives a new session, and why the rule still holds afterwards
- The mappings from sessions to servers and from servers to sessions, and the no-op for an ID that is already active
- A snapshot that includes empty servers and shares no mutable state with the manager
### Part 2 — End sessions and rebalance
Implement `end_session`. Removing a session can break the rule, as `end_session("c")` does in Example 1; restore it by moving sessions, without creating or deleting any. Ending an unknown session does nothing, and an ended ID can be started again, as in Example 2.
```hint Case by case
Removing one session lowers exactly one server's load by one. Ask which loads that server and the other servers could have had just before the call, and in which of those cases the rule breaks.
```
#### What This Part Should Cover
- A complete case analysis of when an end breaks the rule, and how few moves restore it
- Keeping both mappings consistent after a move, so that ending a moved session removes it from its current server
- The unknown-session no-op and clean reuse of an ended ID
### Part 3 — Sublinear updates and stated guarantees
Make `start_session` and `end_session` avoid scanning servers or sessions, reaching O(log S) per update or better. Label each bound as worst-case, average-case or amortized, state any assumption you make about hash-table operations, and explain whether memory grows with the number of past operations.
```hint What the rule buys you
The balancing rule limits which loads can exist at any moment. Consider what that implies for the structure that finds a lightest or a heaviest server.
```
#### What This Part Should Cover
- How a lightest server and a heaviest server are found without a scan
- Per-operation bounds, each labelled worst-case, average-case or amortized, with the hash-table assumption stated
- Whether memory depends only on the current sessions or also on the history of operations
### What a Strong Answer Covers
- Uses the balancing rule to pin down which loads can exist before and after each call
- Correct start, end and snapshot semantics, including both no-ops, ID reuse and snapshot independence
- Rebalancing that moves only as many sessions as the rule requires and keeps every mapping consistent
- Update cost that never scans the servers or the sessions, with guarantees labelled precisely
- Memory bounded by the current state rather than by the number of operations
- Tests that check the stated properties instead of one exact allocation
### Follow-up Questions
- Servers can now be added or removed at runtime. How many sessions must move, and what does each reconfiguration cost?
- Each session now has its own integer weight instead of 1. Can the max-minus-min ≤ 1 rule always be met, and what would you guarantee instead?
- Many threads call the manager at once. What must be atomic, and how do you keep `get_allocation` from blocking updates for O(S + A)?
- Moving a session is expensive in production. How would you choose which session to move, and what does that choice cost per update?
Overview: Implement a session manager that assigns sessions to a fixed set of servers and keeps the highest and lowest server loads within one after every start and end. It tests invariant reasoning, rebalancing by migration, consistent two-way mappings, independent snapshots, and sublinear updates with clearly labelled worst-case, average-case and amortized bounds.