# How to Implement a Custom Load Balancer for CoSky: A Complete Guide

> Learn to implement a custom load balancer for CoSky by extending AbstractLoadBalancer and creating a Chooser. Override CoSky's default strategy with your own selection algorithm. Get the complete guide.

- Repository: [Ahoo Wang/cosky](https://github.com/ahoo-wang/cosky)
- Tags: how-to-guide
- Published: 2026-02-23

---

**To implement a custom load balancer for CoSky, extend the `AbstractLoadBalancer` class, implement the nested `Chooser` interface with your selection algorithm, and register the implementation as a `@Primary` Spring bean to override the default weighted random strategy.**

CoSky provides a pluggable service discovery architecture that separates load-balancing logic from service instance management. Whether you need round-robin, least-connections, or a custom weighted algorithm, you can implement your own strategy by leveraging the `LoadBalancer` interface and the reactive infrastructure provided by the `AbstractLoadBalancer` base class.

## Understand the Core Load Balancer Architecture

Before writing custom code, you need to understand the two core abstractions that power CoSky's load balancing: the `LoadBalancer` interface and the `AbstractLoadBalancer` base class.

### The LoadBalancer Interface

The `LoadBalancer` interface in [`cosky-discovery/src/main/kotlin/me/ahoo/cosky/discovery/loadbalancer/LoadBalancer.kt`](https://github.com/ahoo-wang/cosky/blob/main/cosky-discovery/src/main/kotlin/me/ahoo/cosky/discovery/loadbalancer/LoadBalancer.kt) defines the contract for all balancers. It declares the primary entry point `choose(namespace: String, serviceId: String): Mono<ServiceInstance?>` and contains the nested `Chooser` interface, which encapsulates the actual selection logic. When implementing a custom balancer, you provide a concrete `Chooser` that determines which service instance handles each request.

### The AbstractLoadBalancer Base Class

Rather than implementing `LoadBalancer` directly, extend `AbstractLoadBalancer` from [`cosky-discovery/src/main/kotlin/me/ahoo/cosky/discovery/loadbalancer/AbstractLoadBalancer.kt`](https://github.com/ahoo-wang/cosky/blob/main/cosky-discovery/src/main/kotlin/me/ahoo/cosky/discovery/loadbalancer/AbstractLoadBalancer.kt). This abstract class handles the complex reactive plumbing, including watching for service instance changes via the `InstanceEventListenerContainer` and caching instance lists. It delegates the final selection to your `Chooser` implementation through the abstract method `createChooser(instances: List<ServiceInstance>): C`, where `C` is your chooser type.

Reference implementations such as [`BinaryWeightRandomLoadBalancer.kt`](https://github.com/ahoo-wang/cosky/blob/main/BinaryWeightRandomLoadBalancer.kt) and [`RandomLoadBalancer.kt`](https://github.com/ahoo-wang/cosky/blob/main/RandomLoadBalancer.kt) demonstrate how built-in balancers extend this base class to provide specific algorithms.

## Create Your Custom Load Balancer Implementation

To create a functional custom balancer, you must extend `AbstractLoadBalancer` and provide a thread-safe `Chooser` implementation. Below is a complete round-robin load balancer that cycles through instances in a thread-safe manner.

```kotlin
package me.ahoo.cosky.discovery.loadbalancer

import me.ahoo.cosky.discovery.ServiceInstance
import me.ahoo.cosky.discovery.ServiceDiscovery
import me.ahoo.cosky.discovery.event.InstanceEventListenerContainer
import org.springframework.stereotype.Component
import org.springframework.context.annotation.Primary
import java.util.concurrent.atomic.AtomicInteger

@Component
@Primary
class RoundRobinLoadBalancer(
    serviceDiscovery: ServiceDiscovery,
    instanceEventListenerContainer: InstanceEventListenerContainer
) : AbstractLoadBalancer<RoundRobinLoadBalancer.RoundRobinChooser>(serviceDiscovery, instanceEventListenerContainer) {

    override fun createChooser(instances: List<ServiceInstance>): RoundRobinChooser {
        return RoundRobinChooser(instances)
    }

    class RoundRobinChooser(private val instances: List<ServiceInstance>) : LoadBalancer.Chooser {
        private val idx = AtomicInteger(0)

        override fun choose(): ServiceInstance? {
            if (instances.isEmpty()) {
                return null
            }
            val pos = Math.abs(idx.getAndIncrement() % instances.size)
            return instances[pos]
        }
    }
}

```

**Key implementation details:**

- **`@Component @Primary`**: These annotations ensure Spring injects your balancer wherever a `LoadBalancer` is required, overriding the default bean.
- **`createChooser`**: This method is called whenever the service instance list changes, creating a fresh `Chooser` with the updated instances.
- **`AtomicInteger`**: The round-robin index uses an atomic integer to guarantee thread-safe increments across concurrent requests without locking overhead.
- **`Chooser.choose()`**: This method contains your core algorithm logic and returns the selected `ServiceInstance` or `null` if no instances are available.

You can replace the `choose()` method's logic with any algorithm—such as consistent hashing, least-response-time, or weight-based selection—while maintaining the same structure.

## Register the Custom Balancer in Spring

CoSky's auto-configuration class [`CoSkyDiscoveryAutoConfiguration.kt`](https://github.com/ahoo-wang/cosky/blob/main/CoSkyDiscoveryAutoConfiguration.kt) (located in `cosky-spring-cloud-starter-discovery`) defines the default load balancer using `@ConditionalOnMissingBean`:

```kotlin
@Bean
@ConditionalOnMissingBean
fun coSkyLoadBalancer(
    serviceDiscovery: ServiceDiscovery,
    instanceEventListenerContainer: InstanceEventListenerContainer
): LoadBalancer {
    return BinaryWeightRandomLoadBalancer(serviceDiscovery, instanceEventListenerContainer)
}

```

Because the default bean is annotated with `@ConditionalOnMissingBean`, Spring will only instantiate it if no other `LoadBalancer` bean exists in the context. By marking your custom implementation with `@Primary` (as shown in the previous section), you automatically disable the default `BinaryWeightRandomLoadBalancer` without modifying configuration files.

Alternatively, if you prefer explicit configuration over component scanning, define the bean in a configuration class:

```kotlin
@Configuration
class CustomLoadBalancerConfig {

    @Bean
    @Primary
    fun roundRobinLoadBalancer(
        serviceDiscovery: ServiceDiscovery,
        instanceEventListenerContainer: InstanceEventListenerContainer
    ): LoadBalancer {
        return RoundRobinLoadBalancer(serviceDiscovery, instanceEventListenerContainer)
    }
}

```

For environment-specific selection, add `@ConditionalOnProperty` to toggle between implementations via [`application.yaml`](https://github.com/ahoo-wang/cosky/blob/main/application.yaml):

```kotlin
@Bean
@ConditionalOnProperty(name = ["cosky.load-balancer"], havingValue = "round-robin")
@Primary
fun roundRobinLoadBalancer(
    serviceDiscovery: ServiceDiscovery,
    instanceEventListenerContainer: InstanceEventListenerContainer
): LoadBalancer {
    return RoundRobinLoadBalancer(serviceDiscovery, instanceEventListenerContainer)
}

```

## Verify Integration with the REST API

Once registered, your custom balancer is automatically injected into CoSky's REST layer. The [`ServiceController.kt`](https://github.com/ahoo-wang/cosky/blob/main/ServiceController.kt) in [`cosky-rest-api/src/main/kotlin/me/ahoo/cosky/rest/service/ServiceController.kt`](https://github.com/ahoo-wang/cosky/blob/main/cosky-rest-api/src/main/kotlin/me/ahoo/cosky/rest/service/ServiceController.kt) uses constructor injection to obtain the `LoadBalancer`:

```kotlin
@RestController
class ServiceController(
    private val loadBalancer: LoadBalancer,
    // additional dependencies...
) {
    // Endpoints such as /cosky/instance/choose use the injected balancer
}

```

When you invoke service discovery endpoints, the controller delegates to your custom `choose()` implementation. You can verify the integration by monitoring instance selection patterns or adding logging to your `Chooser.choose()` method.

## Summary

- **Extend `AbstractLoadBalancer`** to inherit reactive instance watching and caching infrastructure from `cosky-discovery`.
- **Implement the `Chooser` interface** to encapsulate your specific selection algorithm (round-robin, weighted, etc.) in a thread-safe manner.
- **Register as `@Primary`** to override the default `BinaryWeightRandomLoadBalancer` defined in [`CoSkyDiscoveryAutoConfiguration.kt`](https://github.com/ahoo-wang/cosky/blob/main/CoSkyDiscoveryAutoConfiguration.kt).
- **Leverage `@ConditionalOnMissingBean`** behavior so your custom implementation automatically disables the default without explicit configuration changes.
- **Verify through `ServiceController`** that REST API calls route through your custom selection logic.

## Frequently Asked Questions

### What interface must I implement to create a custom CoSky load balancer?

You must implement the `LoadBalancer` interface or, more conveniently, extend the `AbstractLoadBalancer` abstract class from the `me.ahoo.cosky.discovery.loadbalancer` package. The abstract class requires you to implement `createChooser()`, which returns an instance of the `LoadBalancer.Chooser` interface containing your `choose()` logic.

### How do I ensure CoSky uses my load balancer instead of the default?

Annotate your implementation with `@Primary` and register it as a Spring bean using `@Component` or an explicit `@Bean` method. Because the default `BinaryWeightRandomLoadBalancer` in [`CoSkyDiscoveryAutoConfiguration.kt`](https://github.com/ahoo-wang/cosky/blob/main/CoSkyDiscoveryAutoConfiguration.kt) is marked with `@ConditionalOnMissingBean`, Spring will inject your primary bean instead of the default whenever your custom implementation is present in the application context.

### Is the custom load balancer thread-safe for high-concurrency scenarios?

Thread safety depends on your `Chooser` implementation. The `AbstractLoadBalancer` recreates the `Chooser` instance whenever the service instance list changes, but multiple threads may concurrently call `choose()` on the same chooser instance. Use thread-safe constructs such as `AtomicInteger` for round-robin counters or concurrent data structures for connection tracking to ensure consistency under load.

### Can I configure which load balancer to use via application properties?

Yes. Use Spring's `@ConditionalOnProperty` annotation on your bean definition to enable or disable specific balancer implementations based on configuration values. For example, you can create multiple balancer beans and activate them conditionally using properties like `cosky.load-balancer=round-robin`, allowing runtime selection without code changes.