loadbalancer

package com.mokaddesh.ratelimtier.service;

import com.mokaddesh.ratelimtier.model.Server;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.locks.ReentrantReadWriteLock;

public class LoadBalancerImpl implements LoadBalancer {

    private final List<Server> servers = new ArrayList<>();
    private final AtomicInteger lastUsedIndex = new AtomicInteger();
    private final ScheduledExecutorService scheduledExecutorService = Executors.newScheduledThreadPool(1);
    private final HealthMonitoring healthMonitoring;
    private final ReentrantReadWriteLock reentrantReadWriteLock = new ReentrantReadWriteLock();

    public LoadBalancerImpl(HealthMonitoring healthMonitoring) {
        this.healthMonitoring = healthMonitoring;
    }

    @Override
    public void addServer(Server server) {
        executeWriteLockOperation(() -> servers.add(server));
    }

    @Override
    public void removeServer(Server server) {
        executeWriteLockOperation(() -> servers.remove(server));
    }

    @Override
    public Server getServer() {

        reentrantReadWriteLock.readLock().lock();
        try {
            if (servers.isEmpty()) {
                throw new RuntimeException("No server");
            }
            int attempt = servers.size();

            while (attempt > 0) {
                int currentNodeIdx = (lastUsedIndex.addAndGet(1)) % servers.size();
                if (servers.get(currentNodeIdx).isActive()) {
                    return servers.get(currentNodeIdx);
                }
                attempt--;
            }

            throw new RuntimeException("No active server");
        } finally {
            reentrantReadWriteLock.readLock().unlock();
        }
    }

    @Override
    public List<Server> getAllServer() {
        return servers;
    }

    @Override
    public void performHealthCheck() {
        scheduledExecutorService.schedule(
                () -> {
                    executeWriteLockOperation(() ->
                            servers.forEach(server ->
                                    server.setActive(healthMonitoring.checkHealth(server)
                    )));
                }, 500, TimeUnit.MILLISECONDS
        );
    }

    private void executeWriteLockOperation(Runnable runnable) {
        reentrantReadWriteLock.writeLock().lock();
        try {
            runnable.run();
        } finally {
            reentrantReadWriteLock.writeLock().unlock();
        }
    }
}