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

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 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. 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 and 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.

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 (located in cosky-spring-cloud-starter-discovery) defines the default load balancer using @ConditionalOnMissingBean:

@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:

@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:

@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 in cosky-rest-api/src/main/kotlin/me/ahoo/cosky/rest/service/ServiceController.kt uses constructor injection to obtain the LoadBalancer:

@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.
• 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 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.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}