io.github.opensabre.rpc.loadbalance.OpensabreLoadBalancer Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of opensabre-starter-rpc Show documentation
The newest version!
package io.github.opensabre.rpc.loadbalance;

import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.loadbalancer.*;
import org.springframework.cloud.loadbalancer.core.NoopServiceInstanceListSupplier;
import org.springframework.cloud.loadbalancer.core.ReactorServiceInstanceLoadBalancer;
import org.springframework.cloud.loadbalancer.core.SelectedInstanceCallback;
import org.springframework.cloud.loadbalancer.core.ServiceInstanceListSupplier;
import reactor.core.publisher.Mono;

import java.util.List;
import java.util.concurrent.atomic.AtomicInteger;


/**
 * Opensabre负载均衡器
 */
@Slf4j
public class OpensabreLoadBalancer implements ReactorServiceInstanceLoadBalancer {

    final AtomicInteger position;

    final String serviceId;

    ObjectProvider serviceInstanceListSupplierProvider;

    /**
     * @param serviceInstanceListSupplierProvider 服务节点
     * @param serviceId                           服务ID，应用名
     */
    public OpensabreLoadBalancer(ObjectProvider serviceInstanceListSupplierProvider, String serviceId) {
        this.serviceId = serviceId;
        this.serviceInstanceListSupplierProvider = serviceInstanceListSupplierProvider;
        this.position = new AtomicInteger(90);
    }

    @Override
    public Mono> choose(Request request) {
        RequestDataContext context = (RequestDataContext) request.getContext();
        RequestData clientRequest = context.getClientRequest();

        ServiceInstanceListSupplier supplier = serviceInstanceListSupplierProvider.getIfAvailable(NoopServiceInstanceListSupplier::new);
        return supplier.get(request).next().map(serviceInstances -> processInstanceResponse(supplier, serviceInstances));
    }

    /**
     * @param supplier         服务实例提供者
     * @param serviceInstances 可用服务实例
     * @return Response ServiceInstance
     */
    private Response processInstanceResponse(ServiceInstanceListSupplier supplier,
                                                              List serviceInstances) {
        Response serviceInstanceResponse = getInstanceResponse(serviceInstances);
        if (supplier instanceof SelectedInstanceCallback && serviceInstanceResponse.hasServer()) {
            ((SelectedInstanceCallback) supplier).selectedServiceInstance(serviceInstanceResponse.getServer());
        }
        return serviceInstanceResponse;
    }

    /**
     * @param instances 可用实例
     * @return Response ServiceInstance
     */
    private Response getInstanceResponse(List instances) {
        if (instances.isEmpty()) {
            log.warn("No servers available for service: " + serviceId);
            return new EmptyResponse();
        }
        //此处编写自己的负载均衡策略
        int pos = this.position.incrementAndGet() & Integer.MAX_VALUE;
        ServiceInstance instance = instances.get(pos % instances.size());
        return new DefaultResponse(instance);
    }
}