服务摘流核心逻辑

当某个服务挂掉后，可以通过下发命令，使某个服务实例摘流，其核心逻辑如下 ：

import com.alibaba.fastjson.JSON;
import com.***.flow.dto.ServiceFlowMqCallBackDTO;
import com.***.advance.flow.dto.ServiceFlowMqExchangeDTO;
import com.***.advance.flow.enums.ServiceStatusEnum;
import com.***.advance.flow.utils.InetUtil;
import com.***.base.json.JsonUtil;
import com.***.base.object.AskBeanUtil;
import com.***.communicate.async2mq.component.ServiceFlowMQProducer;
import com.***.rocketmq.enums.XXXMqEnums;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.apache.curator.x.discovery.ServiceDiscovery;
import org.apache.dubbo.common.extension.ExtensionLoader;
import org.apache.dubbo.common.utils.NetUtils;
import org.apache.dubbo.registry.Registry;
import org.apache.dubbo.registry.RegistryFactory;
import org.apache.dubbo.rpc.model.ApplicationModel;
import org.apache.dubbo.rpc.model.ProviderModel;
import org.apache.dubbo.rpc.model.ServiceRepository;
import org.apache.rocketmq.client.consumer.DefaultMQPushConsumer;
import org.apache.rocketmq.client.consumer.MQPushConsumer;
import org.apache.rocketmq.client.consumer.listener.ConsumeConcurrentlyStatus;
import org.apache.rocketmq.client.consumer.listener.MessageListenerConcurrently;
import org.apache.rocketmq.client.exception.MQClientException;
import org.apache.rocketmq.client.producer.DefaultMQProducer;
import org.apache.rocketmq.common.message.MessageExt;
import org.apache.rocketmq.common.protocol.heartbeat.MessageModel;
import org.apache.rocketmq.remoting.RPCHook;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.stereotype.Component;
import org.springframework.util.CollectionUtils;

import javax.annotation.PostConstruct;
import javax.annotation.Resource;
import java.io.IOException;
import java.net.InetAddress;
import java.util.*;

@Slf4j
@Component
public class ServiceFlowCommandMqListener {

    private final String SERVICE_FLOW_COMMAND_MQ_CONSUMER_BEAN_ID = "serviceFlowCommandMQConsumer";

    private final String SERVICE_FLOW_COMMAND_MQ_PRODUCER = "serviceFlowCommandMQProducer";
    private final String SERVICE_FLOW_COMMAND_MQ_CONSUMER_LISTENER_BEAN_ID = "serviceFlowCommandMQConsumerListener";

    /**
     * 执行后回执消息tag
     */
    public final String SERVICE_FLOW_CALL_BACK_TAG = "service_flow_call_back";
    public final String SERVICE_FLOW_CALL_COMMAND = "service_flow_command";

    @Value("${spring.application.name:}")
    private String APPLICATION_NAME;

    @Autowired(required = false)
    private ServiceDiscovery serviceDiscovery;

    @Autowired(required = false)
    private Map<String, MQPushConsumer> mqConsumerMap;

    @Value("${spring.cloud.zookeeper.discovery.preferIpAddress:true}")
    private Boolean preferIpAddress;

    @Autowired
    private InetUtil inetUtil;

    @PostConstruct
    public void init() {
        getProviderIPTag4DubboApplication();
        getProviderIPTag4HttpApplication();
    }

    public String getProviderIPTag4DubboApplication() {

        InetAddress inetAddress = NetUtils.getLocalAddress();
        if(Objects.isNull(inetAddress)){
            return StringUtils.EMPTY;
        }

        log.info("getProviderIPTag4DubboApplication, ip={}", inetAddress.getHostAddress());
        String ip  = inetAddress.getHostAddress();
        String ipTag = APPLICATION_NAME.replaceAll("-dev|-test|-pre|-prd|-grey", "") + "_" + ip.replace(".", "");
        log.info("getProviderIPTag4DubboApplication, ipTag={}", ipTag);
        return ipTag;
    }

    public String getProviderIPTag4HttpApplication() {
        String ip = getInstanceHost4HttpApplication();
        if(Objects.isNull(ip)){
            log.warn("getProviderIPTag4HttpApplication, get ip failed");
            return StringUtils.EMPTY;
        }

        log.info("getProviderIPTag4HttpApplication, ip={}", ip);
        String ipTag = APPLICATION_NAME.replaceAll("-dev|-test|-pre|-prd|-grey", "") + "_" + ip.replace(".", "");
        log.info("getProviderIPTag4HttpApplication, ipTag={}", ipTag);
        return ipTag;
    }

    private String getInstanceHost4HttpApplication() {
        InetUtil.HostInfo hostInfo = inetUtil.findFirstNonLoopbackHostInfo();
        log.info("getInstanceHost4HttpApplication, hostInfo:{}", JsonUtil.toDefaultString(hostInfo));
        String instanceHost = hostInfo.getHostname();
        String instanceIpAddress = hostInfo.getIpAddress();
        return this.preferIpAddress && org.springframework.util.StringUtils.hasText(instanceIpAddress) ? instanceIpAddress : instanceHost;
    }

    @Bean(name = SERVICE_FLOW_COMMAND_MQ_PRODUCER, initMethod = "start", destroyMethod = "shutdown")
    public ServiceFlowMQProducer getServiceFlowCallBackMQProducer(@Value("${rocketmq.namesrv}") String rocketMqNameSrv, RPCHook rpcHook) throws MQClientException{
        DefaultMQProducer producer = new DefaultMQProducer(DxyMqEnums.SendGroup.SERVICE_FLOW_COMMAND_CALL_BACK_MSG.name(), rpcHook);
        producer.setNamesrvAddr(rocketMqNameSrv);
        producer.setVipChannelEnabled(false);
        return new ServiceFlowMQProducer(producer, SERVICE_FLOW_CALL_BACK_TAG);
    }

    @Resource(name = SERVICE_FLOW_COMMAND_MQ_PRODUCER)
    private ServiceFlowMQProducer serviceFlowMQProducer;


    @Bean(name = SERVICE_FLOW_COMMAND_MQ_CONSUMER_BEAN_ID, initMethod = "start", destroyMethod = "shutdown")
    public DefaultMQPushConsumer serviceFlowCommandMQConsumer(@Qualifier(SERVICE_FLOW_COMMAND_MQ_CONSUMER_LISTENER_BEAN_ID) MessageListenerConcurrently serviceFlowCommandChangeMsgListenerConcurrently,
                                                                 @Value("${rocketmq.namesrv}") String rocketMqNameSrv, RPCHook rpcHook) throws MQClientException {
        DefaultMQPushConsumer consumer = new DefaultMQPushConsumer(rpcHook);
        consumer.setNamesrvAddr(rocketMqNameSrv);
        consumer.setVipChannelEnabled(false);
        consumer.setConsumeTimeout(120_000L);
        consumer.setMessageModel(MessageModel.BROADCASTING);
        //注意如果是消费顺序消息，单线程才能保证消息的顺序到达和顺序消费
        consumer.setConsumeThreadMin(1);
        consumer.setConsumeThreadMax(10);
        consumer.setConsumeMessageBatchMaxSize(1);
        consumer.setInstanceName(UUID.randomUUID().toString());
        consumer.setConsumerGroup(DxyMqEnums.ConsumeGroup.SERVICE_FLOW_COMMAND_MSG.name());
        consumer.subscribe(DxyMqEnums.Topic.SERVICE_FLOW_COMMAND_MSG.name(), SERVICE_FLOW_CALL_COMMAND);
        consumer.registerMessageListener(serviceFlowCommandChangeMsgListenerConcurrently);
        return consumer;
    }

    @Bean(name = SERVICE_FLOW_COMMAND_MQ_CONSUMER_LISTENER_BEAN_ID)
    public MessageListenerConcurrently serviceFlowCommandChangeMsgListenerConcurrently() {
        return (msgs, context) -> {
            if(CollectionUtils.isEmpty(msgs)){
                log.info("serviceFlowCommandChangeMsgListenerConcurrently skipped for msgList is empty.");
                return ConsumeConcurrentlyStatus.CONSUME_SUCCESS;
            }
            for (MessageExt message : msgs) {
                final String msgStr = new String(message.getBody());
                ServiceFlowMqExchangeDTO mqExchangeDTO = JSON.parseObject(msgStr, ServiceFlowMqExchangeDTO.class);
                log.info("serviceFlowCommandChangeMsgListenerConcurrently come, msg:{}, mqExchangeDTO={}", msgStr, mqExchangeDTO);
                if(Objects.isNull(mqExchangeDTO)){
                    continue;
                }
                String uniqueKey = mqExchangeDTO.getApplicationName() + "_" + mqExchangeDTO.getIp().replace(".", "");
                // 不是同一个实例不处理
                if (mqExchangeDTO.getHttpApplication()) {
                    if(!Objects.equals(uniqueKey, getProviderIPTag4HttpApplication())) {
                        continue;
                    }
                } else {
                    if(!Objects.equals(uniqueKey, getProviderIPTag4DubboApplication())) {
                        continue;
                    }
                }

                if (serviceDiscovery != null) {
                    httpHandle(mqExchangeDTO.getStatus());
                }
                dubboHandle(mqExchangeDTO.getStatus());
                if (mqExchangeDTO.getHandleMq()) {
                    mqHandle(mqExchangeDTO.getStatus());
                }
                ServiceFlowMqCallBackDTO mqCallBackDTO = AskBeanUtil.copyProperties(mqExchangeDTO, new ServiceFlowMqCallBackDTO(), false);
                mqCallBackDTO.setExecuteStatus(1);

                serviceFlowMQProducer.sendFlowBackInfo(DxyMqEnums.Topic.SERVICE_FLOW_COMMAND_MSG, SERVICE_FLOW_CALL_BACK_TAG, JSON.toJSONString(mqCallBackDTO));
            }
            return ConsumeConcurrentlyStatus.CONSUME_SUCCESS;
        };
    }



    void httpHandle(Integer status) {
        if (Objects.equals(status, ServiceStatusEnum.ON_LINE.getCode())){
            //http服务摘流
            try {
                log.warn("http service register started");
                serviceDiscovery.start();
                log.warn("http service register finished");
            } catch (Exception e) {
                log.warn("http service register exception", e);
            }
        } else if (Objects.equals(status, ServiceStatusEnum.OFF_LINE.getCode())){
            //http服务恢复流量
            try {
                log.warn("http service unregister started");
                serviceDiscovery.close();
                log.warn("http service unregister started");
            } catch (IOException e) {
                log.warn("http service unregister exception", e);
            }
        }
    }

    void dubboHandle(Integer status){
        RegistryFactory registryFactory = ExtensionLoader.getExtensionLoader(RegistryFactory.class).getAdaptiveExtension();
        if (registryFactory == null) {
            return;
        }
        ServiceRepository serviceRepository = ApplicationModel.getServiceRepository();
        if (serviceRepository == null) {
            return;
        }
        Collection<ProviderModel> providerModelList = serviceRepository.getExportedServices();
        if (CollectionUtils.isEmpty(providerModelList)) {
            return;
        }
        if(Objects.equals(status, ServiceStatusEnum.ON_LINE.getCode())){
            log.warn("Dubbo service register started");
            for (ProviderModel providerModel : providerModelList) {
                List<ProviderModel.RegisterStatedURL> statedUrls = providerModel.getStatedUrl();
                for (ProviderModel.RegisterStatedURL statedURL : statedUrls) {
                    if (!statedURL.isRegistered()) {
                        Registry registry = registryFactory.getRegistry(statedURL.getRegistryUrl());
                        registry.register(statedURL.getProviderUrl());
                        statedURL.setRegistered(true);
                    }
                }
            }
            log.warn("Dubbo service register finished");

        }
        if(Objects.equals(status, ServiceStatusEnum.OFF_LINE.getCode())) {
            log.warn("Dubbo service unregister started");
            for (ProviderModel providerModel : providerModelList) {
                List<ProviderModel.RegisterStatedURL> statedUrls = providerModel.getStatedUrl();
                for (ProviderModel.RegisterStatedURL statedURL : statedUrls) {
                    if (statedURL.isRegistered()) {
                        Registry registry = registryFactory.getRegistry(statedURL.getRegistryUrl());
                        registry.unregister(statedURL.getProviderUrl());
                        statedURL.setRegistered(false);
                    }
                }
            }
            log.warn("Dubbo service unregister finished");
        }

    }

    void mqHandle(Integer status){
        if (CollectionUtils.isEmpty(mqConsumerMap)) {
            log.warn("mqHandle skipped for mqConsumerMap is empty");
            return;
        }
        if(Objects.equals(status, ServiceStatusEnum.ON_LINE.getCode())){
            //RocketMQ Consumer 摘流
            log.warn("RocketMQ Consumer suspend started");
            mqConsumerMap.forEach((name, consumer) -> {
                DefaultMQPushConsumer defaultMQPushConsumer = (DefaultMQPushConsumer) consumer;
                if (!defaultMQPushConsumer.getConsumerGroup().equals(DxyMqEnums.ConsumeGroup.SERVICE_FLOW_COMMAND_MSG.name())) {
                    consumer.suspend();
                }
            });
            log.warn("RocketMQ Consumer suspend finished");
        }
        if(Objects.equals(status, ServiceStatusEnum.OFF_LINE.getCode())){
            //RocketMQ Consumer 恢复流量
            log.warn("RocketMQ Consumer resume started");
            mqConsumerMap.forEach((name, consumer) -> {
                DefaultMQPushConsumer defaultMQPushConsumer = (DefaultMQPushConsumer) consumer;
                if (!defaultMQPushConsumer.getConsumerGroup().equals(DxyMqEnums.ConsumeGroup.SERVICE_FLOW_COMMAND_MSG.name())) {
                    consumer.resume();
                }
            });
            log.warn("RocketMQ Consumer resume finished");
        }

    }
}