Skip to content

Implement your own load balancer mechanism

Stork is extensible, and you can implement your service selection (load-balancer) mechanism.

Dependencies

To implement your Load Balancer Provider, make sure your project depends on Core and Configuration Generator. The former brings classes necessary to implement custom load balancer, the latter contains an annotation processor that generates classes needed by Stork.

<dependency>
    <groupI>io.smallrye.stork</groupI>
    <artifactId>stork-core</artifactId>
    <version>1.0.0</version>
</dependency>
<dependency>
    <groupId>io.smallrye.stork</groupId>
    <artifactId>stork-configuration-generator</artifactId>
    <scope>provided</scope>
    <!-- provided scope is sufficient for the annotation processor -->
    <version>1.0.0</version>
</dependency>

Implementing a load balancer provider

Load balancer implementation consists of three elements:

  • LoadBalancer which is responsible for selecting service instances for a single Stork service
  • LoadBalancerProvider which creates instances of LoadBalancer for a given load balancer type
  • LoadBalancerProviderConfiguration which is a configuration for the load balancer

A type, for example, acme, identifies each provider. This type is used in the configuration to reference the provider:

stork.my-service.load-balancer=acme

A LoadBalancerProvider implementation needs to be annotated with @LoadBalancerType that defines the type. Any configuration properties that the provider expects should be defined with @LoadBalancerAttribute annotations placed on the provider.

A load balancer provider class should look as follows: 

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
package examples;

import io.smallrye.stork.api.LoadBalancer;
import io.smallrye.stork.api.ServiceDiscovery;
import io.smallrye.stork.api.config.LoadBalancerAttribute;
import io.smallrye.stork.api.config.LoadBalancerType;
import io.smallrye.stork.spi.LoadBalancerProvider;

@LoadBalancerType("acme")
@LoadBalancerAttribute(name = "my-attribute",
        description = "Attribute that alters the behavior of the LoadBalancer")
public class AcmeLoadBalancerProvider implements
        LoadBalancerProvider<AcmeLoadBalancerProviderConfiguration> {

    @Override
    public LoadBalancer createLoadBalancer(AcmeLoadBalancerProviderConfiguration config,
                                           ServiceDiscovery serviceDiscovery) {
        return new AcmeLoadBalancer(config);
    }
}

Note, that the LoadBalancerProvider interface takes a configuration class as a parameter. This configuration class is generated automatically by the Configuration Generator. Its name is created by appending Configuration to the name of the provider class.

The next step is to implement the LoadBalancer interface:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
package examples;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Random;

import io.smallrye.stork.api.LoadBalancer;
import io.smallrye.stork.api.NoServiceInstanceFoundException;
import io.smallrye.stork.api.ServiceInstance;

public class AcmeLoadBalancer implements LoadBalancer {

    private final Random random;

    public AcmeLoadBalancer(AcmeLoadBalancerProviderConfiguration config) {
        random = new Random();
    }

    @Override
    public ServiceInstance selectServiceInstance(Collection<ServiceInstance> serviceInstances) {
        if (serviceInstances.isEmpty()) {
            throw new NoServiceInstanceFoundException("No services found.");
        }
        int index = random.nextInt(serviceInstances.size());
        return new ArrayList<>(serviceInstances).get(index);
    }
}

This implementation is simplistic and just picks a random instance from the received list.

Using your load balancer

In the project using it, don't forget to add the dependency on the module providing your implementation. Then, in the configuration, just add:

stork.my-service.service-discovery=...
stork.my-service.load-balancer=acme

Then, Stork will use your implementation to select the my-service service instance.