im grpc推送ws网关

package com.kongjs.im.push.grpc.service;

import com.alibaba.fastjson2.JSON;
import com.alibaba.fastjson2.JSONArray;
import com.alibaba.fastjson2.JSONObject;
import com.kongjs.im.grpc.proto.*;
import io.grpc.stub.StreamObserver;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.redisson.api.RBucket;
import org.redisson.api.RScoredSortedSet;
import org.redisson.api.RedissonClient;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.discovery.DiscoveryClient;
import org.springframework.grpc.client.GrpcChannelFactory;
import org.springframework.stereotype.Service;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;

import java.util.*;
import java.util.stream.Collectors;

@Slf4j
@Service
public class GrpcServerPushMessageService extends PushMessageServiceGrpc.PushMessageServiceImplBase {

    private static final String USER_ONLINE = "im:online:user:{%s}";
    private static final String GROUP_MEMBER = "im:group:member:%s";

    @Resource
    private DiscoveryClient discoveryClient;
    @Resource
    private RedissonClient redissonClient;
    @Resource
    private GrpcChannelFactory grpcChannelFactory;

    @Override
    public void pushMessage(MessagePushReq request, StreamObserver<MessagePushResp> responseObserver) {
        try {
            ChatMessageData data = request.getData();
            ChatType chatType = data.getChatType();
            String to = data.getTo();
            if (chatType == ChatType.SYSTEM) {
                Set<String> gatewayIds = discoveryClient.getInstances("im-ws")
                        .stream()
                        .map(ServiceInstance::getHost)
                        .collect(Collectors.toSet());
                for (String gatewayId : gatewayIds) {
                    String url = gatewayId + ":11001";
                    WebSocketResult result = wsStub(url).forwardMessage(data);
                    if (ResultCode.SUCCESS.equals(result.getCode())) {
                        String data1 = result.getData();
                        JSONObject object = JSON.parseObject(data1);
                        log.info(object.toJSONString());
                    }
                }
            } else if (ChatType.SINGLE.equals(chatType)) {
                String userOnlineKey = String.format(USER_ONLINE, to);
                RBucket<String> bucket = redissonClient.getBucket(userOnlineKey);
                String gatewayId = bucket.get();
                String url = gatewayId + ":11001";
                WebSocketResult result = wsStub(url).forwardMessage(data);
                if (ResultCode.SUCCESS.equals(result.getCode())) {
                    String data1 = result.getData();
                    JSONObject object = JSON.parseObject(data1);
                    log.info(object.toJSONString());
                }
            } else if (ChatType.GROUP.equals(chatType)) {
                String redisKey = String.format(GROUP_MEMBER,to);
                RScoredSortedSet<Long> zSet = redissonClient.getScoredSortedSet(redisKey);
                int size = zSet.size();
                if (size <= 500) {
                    Collection<Long> memberIds = zSet.readAll();
                    if (CollectionUtils.isEmpty(memberIds)) {
                        return;
                    }
                    push(memberIds, data);
                } else {
                    int pageSize = 200;
                    int offset = 0;
                    do {
                        Collection<Long> memberIds = zSet.valueRangeReversed(offset, offset + pageSize - 1);
                        if (CollectionUtils.isEmpty(memberIds)) {
                            break;
                        }
                        push(memberIds, data);
                        offset += pageSize;
                    } while (true);
                }
            }
            responseObserver.onNext(MessagePushResp.newBuilder().setCode(ResultCode.SUCCESS).build());
        } catch (Exception e) {
            responseObserver.onNext(MessagePushResp.newBuilder().setCode(ResultCode.ERROR).setMsg(e.getMessage()).build());
        } finally {
            responseObserver.onCompleted();
        }
    }

    private WebSockerMessageServiceGrpc.WebSockerMessageServiceBlockingStub wsStub(String target) {
        return WebSockerMessageServiceGrpc.newBlockingStub(grpcChannelFactory.createChannel(target));
    }

    public void push(Collection<Long> memberIds, ChatMessageData data) {
        Set<String> keys = memberIds.stream().map(s -> String.format(USER_ONLINE, s)).collect(Collectors.toSet());
        if (CollectionUtils.isEmpty(keys)) {
            return;
        }
        Map<String, String> gatewayMap = redissonClient.getBuckets().get(keys.toArray(new String[0]));
        Map<String, Set<String>> gatewayUsers = new HashMap<>();
        for (Long userId : memberIds) {
            String key = String.format(USER_ONLINE, userId);
            String gatewayId = gatewayMap.get(key);
            if (StringUtils.hasText(gatewayId)) {
                continue;
            }
            // 同一个网关放一起
            gatewayUsers.computeIfAbsent(gatewayId, _ -> new HashSet<>())
                    .add(userId.toString());
        }
        for (Map.Entry<String, Set<String>> entry : gatewayUsers.entrySet()) {
            String gatewayId = entry.getKey();
            Set<String> userIds = entry.getValue();
            String pushUrl = gatewayId + ":11001";
            // 一次请求 → 推送给该网关下所有用户
            ChatMessageData.Builder builder = ChatMessageData.newBuilder().mergeFrom(data)
                    .addAllTos(userIds);
            WebSocketResult result = wsStub(pushUrl).forwardMessage(builder.build());
            if (ResultCode.SUCCESS.equals(result.getCode())) {
                String data1 = result.getData();
                JSONArray array = JSON.parseArray(data1);
                log.info(array.toJSONString());
            }
        }
    }
}