自定义spring gateway负载均衡

自定义一个规则：

import com.netflix.client.config.IClientConfig;
import com.netflix.loadbalancer.AbstractLoadBalancerRule;
import com.netflix.loadbalancer.Server;

import java.util.List;

public class ConsistenceHashRule extends AbstractLoadBalancerRule {
    @Override
    public void initWithNiwsConfig(IClientConfig iClientConfig) {
        System.out.println(iClientConfig);
    }

    @Override
    public Server choose(Object o) {
        List servers = this.getLoadBalancer().getReachableServers();
        return servers.get(1);

    }
}

重写：LoadBalancerClientFilter的实现

其实，基本上不用重写filter方法，但是一般都要进行简单的权限验证或者添加自己的配置，以及添加header等信息，看自己的需求吧

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.loadbalancer.LoadBalancerClient;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.LoadBalancerClientFilter;
import org.springframework.cloud.gateway.support.NotFoundException;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;

import java.net.URI;
import java.util.Map;

import static org.springframework.cloud.gateway.support.ServerWebExchangeUtils.*;

public class WebSocketLBGatewayFilter extends LoadBalancerClientFilter {
    private static final Log log = LogFactory.getLog(WebSocketLBGatewayFilter.class);

    private WSLoadBalancerProperties properties = null;

    public WebSocketLBGatewayFilter(LoadBalancerClient loadBalancer,
                                    WSLoadBalancerProperties properties) {
        super(loadBalancer, properties);
        this.properties = properties;
    }


    @Override
    public Mono filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        URI url = exchange.getAttribute(GATEWAY_REQUEST_URL_ATTR);
        String schemePrefix = exchange.getAttribute(GATEWAY_SCHEME_PREFIX_ATTR);
        String SCHEME_PREFIX = this.properties.getSchemePrefix()==null?"lb":this.properties.getSchemePrefix();
        //配置的时候，需要用lb开头的原因：如果不使用lb开头，则不进行负载均衡
        if (url == null || (!SCHEME_PREFIX.equals(url.getScheme()) && !SCHEME_PREFIX.equals(schemePrefix))) {
            return chain.filter(exchange);
        }
        // preserve the original url
        addOriginalRequestUrl(exchange, url);

        log.trace("LoadBalancerClientFilter url before: " + url);

        final ServiceInstance instance = this.choose(exchange);

        if (instance == null) {
            throw NotFoundException.create(properties.isUse404(),
                    "Unable to find instance for " + url.getHost());
        }

        URI uri = exchange.getRequest().getURI();

        // if the `lb:` mechanism was used, use `` as the default,
        // if the loadbalancer doesn't provide one.
        String overrideScheme = instance.isSecure() ? "wss" : "ws";
        if (schemePrefix != null) {
            overrideScheme = url.getScheme();
        }

        URI requestUrl = loadBalancer.reconstructURI(
                new WSServiceInstance(overrideScheme, instance), uri);

        log.trace("LoadBalancerClientFilter url chosen: " + requestUrl);
        exchange.getAttributes().put(GATEWAY_REQUEST_URL_ATTR, requestUrl);

        return chain.filter(exchange);
    }

    protected ServiceInstance choose(ServerWebExchange exchange) {
        URI uri =  exchange.getAttribute(GATEWAY_REQUEST_URL_ATTR);
        String serviceId = uri.getHost();
        return loadBalancer.choose(serviceId);
    }



    class WSServiceInstance implements ServiceInstance {
        private String schema;
        private final ServiceInstance instance;

        public WSServiceInstance(String schema, ServiceInstance instance) {
            this.schema = schema;
            this.instance = instance;
        }


        @Override
        public String getServiceId() {
            return instance.getServiceId();
        }

        @Override
        public String getHost() {
            return instance.getHost();
        }

        @Override
        public int getPort() {
            return instance.getPort();
        }

        @Override
        public boolean isSecure() {
            return instance.isSecure();
        }

        @Override
        public URI getUri() {
            return instance.getUri();
        }

        @Override
        public Map getMetadata() {
            return instance.getMetadata();
        }
    }

我想添加自己的一些负载均衡的配置，所以继承了一下WSLoadBalancerProperties这个基类：

import lombok.Data;
import org.springframework.cloud.gateway.config.LoadBalancerProperties;

@Data
public class WSLoadBalancerProperties extends LoadBalancerProperties {
    private String schemePrefix ;

}

 

在配置中添加：

spring:
  cloud:
    gateway:
      discovery:
        locator:
          lowerCaseServiceId: true
          enabled: true
      loadbalancer:
        schemePrefix: lb
      default-filters:
      routes:
      - id: terminal-websocket
        uri: lb:ws://terminal-server
        predicates:
        - Path=/websocket/**
        filters:
        - StripPrefix=1

terminal-server:
  robbin:
    NFLoadBalancerRuleClassName:  cn.com.lenovo.ai.dm.www.gateway.rule.ConsistenceHashRule

这个配置的意思是，名叫terminal-server服务使用ConsistenceHashRule这个负载策略（其他服务依然使用默认的负载策略，spring cloud ribbon提供的默认负载策略是这个类com.netflix.loadbalancer.ZoneAvoidanceRule）。

注入上面的所有的实现：

import cn.com.lenovo.ai.dm.www.gateway.filter.WSLoadBalancerProperties;
import cn.com.lenovo.ai.dm.www.gateway.filter.WebSocketLBGatewayFilter;
import org.springframework.cloud.client.loadbalancer.LoadBalancerClient;
import org.springframework.cloud.gateway.filter.LoadBalancerClientFilter;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Primary;

@Configuration
public class LBConfiguration {


    @Bean
    public LoadBalancerClientFilter webSocketLBGatewayFilter(LoadBalancerClient loadBalancerClient, WSLoadBalancerProperties wsLoadBalancerProperties) {
        return new WebSocketLBGatewayFilter(loadBalancerClient, wsLoadBalancerProperties);
    }

    @Bean
    @Primary
    WSLoadBalancerProperties wsLoadBalancerProperties() {
        return new WSLoadBalancerProperties();
    }


//    @Bean
//    public WebSocketGatewayFilterFactory authorizeGatewayFilterFactory(LoadBalancerClient loadBalancerClient,
//                                                                       LoadBalancerProperties properties) {
//        return new WebSocketGatewayFilterFactory(loadBalancerClient, properties);
//    }

}

前端请求websocket的地址：ws://localhost:11000/websocket/ws/${userId}

${userId}可以随自己的需求随意替换；