DEV Community

Sandeep Vishnu
Sandeep Vishnu

Posted on

πŸ—οΈ Building a Simple Load Balancer with Spring Boot: A Step-by-Step Guide

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:

  1. Load Balancer ServiceΒ βš–οΈ: The core component that handles request distribution
  2. API ServiceΒ πŸ”Œ: Multiple instances of a demo service that will receive the distributed traffic
  3. 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:

High-level view of our system architecture

The system works as follows:

  1. Clients send requests to the load balancer on port 8080 πŸ“‘
  2. The load balancer maintains a pool of healthy API services 🏊
  3. API services send heartbeat messages every 5 seconds to register themselves πŸ’—
  4. The Health Check Service monitors service health and removes unresponsive instances πŸ₯
  5. The Load Balancer Service distributes incoming requests across healthy instances using round-robin πŸ”„
  6. 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
) {}
Enter fullscreen mode Exit fullscreen mode

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
) {}
Enter fullscreen mode Exit fullscreen mode

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());
    }
}
Enter fullscreen mode Exit fullscreen mode

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()));
    }
}
Enter fullscreen mode Exit fullscreen mode

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());
    }
}
Enter fullscreen mode Exit fullscreen mode

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());
        }
    }
}

Enter fullscreen mode Exit fullscreen mode

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
Enter fullscreen mode Exit fullscreen mode

Key Features ✨

  1. Round-Robin Load Balancing πŸ”„: Requests are distributed evenly across healthy service instances
  2. Health Checking πŸ₯: Regular heartbeat monitoring ensures only healthy instances receive traffic
  3. Dynamic Service Registration πŸ“: Services can join or leave the cluster at any time
  4. Request Forwarding πŸ“¨: All HTTP methods and headers are properly forwarded
  5. Docker Support 🐳: Easy deployment with Docker Compose
  6. Modular Design 🧩: Clean separation of concerns with distinct modules

Testing the Load Balancer πŸ§ͺ

To test the load balancer:

  1. Start the system using Docker Compose:
docker-compose up --build
Enter fullscreen mode Exit fullscreen mode
  1. Send requests to the load balancer (port 8080):
curl http://localhost:8080/api/demo
Enter fullscreen mode Exit fullscreen mode

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
Enter fullscreen mode Exit fullscreen mode

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. πŸŽ“

Top comments (0)