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>
    <groupId>io.smallrye.stork</groupId>
    <artifactId>stork-core</artifactId>
    <version>2.6.3</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>2.6.3</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,
• $typeConfiguration which is a configuration for the load balancer. This class is automatically generated.

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

stork standalonestork in quarkus

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

quarkus.stork.my-service.load-balancer.type=acme-load-balancer 

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. Optionally, you can also add @ApplicationScoped annotation in order to provide the load balancer implementation as CDI bean.

A load balancer provider class should look as follows:

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;
import jakarta.enterprise.context.ApplicationScoped;

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

    @Override
    public LoadBalancer createLoadBalancer(AcmeLoadBalancerConfiguration 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 load balancer type, like AcmeLoadBalancerConfiguration.

The next step is to implement the LoadBalancer interface.

The essence of load balancers’ work happens in the selectServiceInstance method. The method returns a single ServiceInstance from a collection.

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

    @Override
    public boolean requiresStrictRecording() {
        return false;
    }
}

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

Some load balancers make the pick based on statistics such as calls in progress or response times, or amount of errors of a service instance. To collect this information in your load balancer, you can wrap the selected service instance into ServiceInstanceWithStatGathering.

Load balancers based on statistics often expect that an operation using a selected service instance is marked as started before the next selection. By default, Stork assumes that a LoadBalancer requires this and guards the calls accordingly. If this is not the case for your implementation, override the requiresStrictRecording() method to return false.

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 standalonestork in quarkus

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

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

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

Using your load balancer using the programmatic API

When building your load balancer project, the configuration generator creates a configuration class. This class can be used to configure your load balancer using the Stork programmatic API.

package examples;

import io.smallrye.mutiny.Uni;
import io.smallrye.stork.api.ServiceDefinition;
import io.smallrye.stork.api.ServiceInstance;
import io.smallrye.stork.api.StorkServiceRegistry;
import io.smallrye.stork.servicediscovery.staticlist.StaticConfiguration;
import io.smallrye.stork.servicediscovery.staticlist.StaticRegistrarConfiguration;

public class AcmeSelectorApiUsage {

    public void example(StorkServiceRegistry stork) {
        String list = "localhost:8080, localhost:8081";
        stork.defineIfAbsent("my-service", ServiceDefinition.of(
                new StaticConfiguration().withAddressList(list),
                new AcmeLoadBalancerConfiguration().withMyAttribute("my-value"),new StaticRegistrarConfiguration())
        );

        Uni<ServiceInstance> uni = stork.getService("my-service").selectInstance();
    }

}

Remember that attributes, like my-attribute, are declared using the @LoadBalancerAttribute annotation on the LoadBalancerProvider implementation.