🏗️ Building a Simple Load Balancer with Spring Boot: A Step-by-Step Guide

Sandeep Vishnu — Tue, 05 Nov 2024 01:35:42 +0000

In today's microservices architecture, load balancers play a crucial role in distributing traffic across multiple service instances. This article will walk you through building a simple yet functional load balancer using Spring Boot, complete with health checking and round-robin load distribution.

Project Overview 📋

Our load balancer implementation consists of three main components:

Load Balancer Service ⚖️: The core component that handles request distribution
API Service 🔌: Multiple instances of a demo service that will receive the distributed traffic
Common Module 📦: Shared DTOs and utilities

The complete project uses Spring Boot 3.2.0 and Java 21 ☕, showcasing modern Java features and enterprise-grade patterns.

Architecture 🏗️

Here's a high-level view of our system architecture:

The system works as follows:

Clients send requests to the load balancer on port 8080 📡
The load balancer maintains a pool of healthy API services 🏊
API services send heartbeat messages every 5 seconds to register themselves 💗
The Health Check Service monitors service health and removes unresponsive instances 🏥
The Load Balancer Service distributes incoming requests across healthy instances using round-robin 🔄
Each API service runs on a different port (8081-8085) and processes the forwarded requests 🔌

Let's dive into each component in detail.

Implementation 💻

1. Common DTOs 📝

First, let's look at our shared data structures. These are used for communication between services:

// HeartbeatRequest.java
public record HeartbeatRequest(
        String serviceId,
        String host,
        int port,
        String status,
        long timestamp
) {}

// HeartbeatResponse.java
public record HeartbeatResponse(
        boolean acknowledged,
        String message,
        long timestamp
) {}

2. Service Node Model 🏢

The load balancer keeps track of service instances using the ServiceNode record:

public record ServiceNode(
        String serviceId,
        String host,
        int port,
        boolean healthy,
        Instant lastHeartbeat
) {}

3. Load Balancer Service ⚖️

The core load balancing logic is implemented in LoadBalancerService:

@Service
public class LoadBalancerService {
    private final ConcurrentHashMap<String, ServiceNode> serviceNodes = new ConcurrentHashMap<>();
    private final AtomicInteger currentNodeIndex = new AtomicInteger(0);

    public void registerNode(ServiceNode node) {
        serviceNodes.put(node.serviceId(), node);
    }

    public void removeNode(String serviceId) {
        serviceNodes.remove(serviceId);
    }

    public ServiceNode getNextAvailableNode() {
        List<ServiceNode> healthyNodes = serviceNodes.values().stream()
                .filter(ServiceNode::healthy)
                .toList();

        if (healthyNodes.isEmpty()) {
            throw new IllegalStateException("No healthy nodes available");
        }

        int index = currentNodeIndex.getAndIncrement() % healthyNodes.size();
        return healthyNodes.get(index);
    }

    public List<ServiceNode> getAllNodes() {
        return new ArrayList<>(serviceNodes.values());
    }
}

4. Health Check Service 🏥

The HealthCheckService manages service registration and health monitoring:

@Service
@Slf4j
public class HealthCheckService {
    private final LoadBalancerService loadBalancerService;
    private static final long HEALTH_CHECK_TIMEOUT_SECONDS = 30;

    public HealthCheckService(LoadBalancerService loadBalancerService) {
        this.loadBalancerService = loadBalancerService;
    }

    public HeartbeatResponse processHeartbeat(HeartbeatRequest request) {
        ServiceNode node = new ServiceNode(
                request.serviceId(),
                request.host(),
                request.port(),
                true,
                Instant.now()
        );
        loadBalancerService.registerNode(node);

        return new HeartbeatResponse(true, "Heartbeat acknowledged",
                Instant.now().toEpochMilli());
    }

    @Scheduled(fixedRate = 10000)// Check every 10 seconds
    public void checkNodeHealth() {
        Instant threshold = Instant.now().minus(HEALTH_CHECK_TIMEOUT_SECONDS,
                ChronoUnit.SECONDS);

        loadBalancerService.getAllNodes().stream()
                .filter(node -> node.lastHeartbeat().isBefore(threshold))
                .forEach(node -> loadBalancerService.removeNode(node.serviceId()));
    }
}

5. Proxy Controller 🔄

The ProxyController handles incoming requests and forwards them to the appropriate service:

@Slf4j
@RestController
public class ProxyController {
    private final LoadBalancerService loadBalancerService;
    private final HealthCheckService healthCheckService;
    private final RestTemplate restTemplate;

    @PostMapping("/heartbeat")
    public HeartbeatResponse handleHeartbeat(@RequestBody HeartbeatRequest request) {
        return healthCheckService.processHeartbeat(request);
    }

    @RequestMapping(value = "/**")
    public ResponseEntity<?> proxyRequest(HttpServletRequest request)
            throws URISyntaxException, IOException {
        var node = loadBalancerService.getNextAvailableNode();
        String targetUrl = String.format("http://%s:%d%s",
                node.host(),
                node.port(),
                request.getRequestURI()
        );

// Copy headers
        HttpHeaders headers = new HttpHeaders();
        Enumeration<String> headerNames = request.getHeaderNames();
        while (headerNames.hasMoreElements()) {
            String headerName = headerNames.nextElement();
            headers.addAll(headerName,
                    Collections.list(request.getHeaders(headerName)));
        }

// Forward the request
        ResponseEntity<String> response = restTemplate.exchange(
                new URI(targetUrl),
                HttpMethod.valueOf(request.getMethod()),
                new HttpEntity<>(StreamUtils.copyToByteArray(
                        request.getInputStream()), headers),
                String.class
        );

        return new ResponseEntity<>(response.getBody(),
                response.getHeaders(),
                response.getStatusCode());
    }
}

6. API Service Implementation 🔌

The API service includes a heartbeat configuration to register with the load balancer:

@Slf4j
@Component
public class HeartbeatConfig {
    private final RestTemplate restTemplate;
    private final String serviceId = UUID.randomUUID().toString();

    @Value("${server.port}")
    private int serverPort;

    @Value("${loadbalancer.url}")
    private String loadBalancerUrl;

    @Scheduled(fixedRate = 5000)// Send heartbeat every 5 seconds
    public void sendHeartbeat() {
        try {
            String hostname = InetAddress.getLocalHost().getHostName();

            var request = new HeartbeatRequest(
                    serviceId,
                    hostname,
                    serverPort,
                    "UP",
                    Instant.now().toEpochMilli()
            );

            restTemplate.postForObject(
                    loadBalancerUrl + "/heartbeat",
                    request,
                    void.class
            );

            log.info("Heartbeat sent successfully to {}", loadBalancerUrl);
        } catch (Exception e) {
            log.error("Failed to send heartbeat: {}", e.getMessage());
        }
    }
}

Deployment with Docker Compose 🐳

The project includes Docker support for easy deployment. Here's a snippet from the docker-compose.yml:

services:
  load-balancer:
    build:
      context: .
      dockerfile: load-balancer/Dockerfile
    ports:
      - "8080:8080"
    networks:
      - app-network

  api-service-1:
    build:
      context: .
      dockerfile: api-service/Dockerfile
    environment:
      - SERVER_PORT=8081
      - LOADBALANCER_URL=http://load-balancer:8080
    networks:
      - app-network

  api-service-2:
    build:
      context: .
      dockerfile: api-service/Dockerfile
    environment:
      - SERVER_PORT=8082
      - LOADBALANCER_URL=http://load-balancer:8080
    networks:
      - app-network

networks:
  app-network:
    driver: bridge

Key Features ✨

Round-Robin Load Balancing 🔄: Requests are distributed evenly across healthy service instances
Health Checking 🏥: Regular heartbeat monitoring ensures only healthy instances receive traffic
Dynamic Service Registration 📝: Services can join or leave the cluster at any time
Request Forwarding 📨: All HTTP methods and headers are properly forwarded
Docker Support 🐳: Easy deployment with Docker Compose
Modular Design 🧩: Clean separation of concerns with distinct modules

Testing the Load Balancer 🧪

To test the load balancer:

Start the system using Docker Compose:

docker-compose up --build

Send requests to the load balancer (port 8080):

curl http://localhost:8080/api/demo

You should see responses from different service instances as the load balancer distributes the requests. 🔄

Monitoring and Metrics 📊

The application includes Spring Boot Actuator endpoints for monitoring:

management:
  endpoints:
    web:
      exposure:
        include: health,info,metrics
  endpoint:
    health:
      show-details: always

Conclusion 🎯

This implementation demonstrates a simple but functional load balancer using Spring Boot. While it may not have all the features of production-grade load balancers like Nginx or HAProxy, it serves as an excellent learning tool and could be extended with additional features like:

Weighted round-robin ⚖️
Least connections algorithm 🔗
Sticky sessions 🍪
Circuit breakers 🔌
Rate limiting 🚦

For reference, the entire code implementation can also be found at this Github Repository: https://github.com/sandeepkv93/SimpleLoadBalancer 📚

Remember that in production environments, you might want to use battle-tested solutions like Nginx, HAProxy, or cloud provider load balancers. However, understanding how load balancers work under the hood is valuable knowledge for any software engineer. 🎓

DEV Community: Sandeep Vishnu