实时获取 springboot-websocket-session

WebSocketSessionContainer

import org.springframework.stereotype.Component;
import org.springframework.util.ClassUtils;
import org.springframework.util.ReflectionUtils;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.messaging.SubProtocolWebSocketHandler;

import javax.inject.Inject;
import java.lang.reflect.Field;
import java.lang.reflect.ParameterizedType;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.function.Function;
import java.util.stream.Collectors;

@Component
public class WebSocketSessionContainer {
    private final Map webSocketSessionMap;

    private final SubProtocolWebSocketHandler        websocketHandler;
    private final Field                              field_session;
    private final Function transferFn;

    @Inject
    public WebSocketSessionContainer(WebSocketHandler websocketHandler) throws ClassNotFoundException {
        this.websocketHandler = SubProtocolWebSocketHandler.class.cast(websocketHandler);

        Field field_sessions = ReflectionUtils.findField(SubProtocolWebSocketHandler.class, "sessions");
        ReflectionUtils.makeAccessible(field_sessions);
        Object idToSessionHolder = ReflectionUtils.getField(field_sessions, websocketHandler);
        this.webSocketSessionMap = idToSessionHolder != null ? (Map) idToSessionHolder : Map.of();

        String   typeName = ((ParameterizedType) field_sessions.getGenericType()).getActualTypeArguments()[1].getTypeName();
        Class clazz    = ClassUtils.forName(typeName, null);  //WebSocketSessionHolder.class
        this.field_session = ReflectionUtils.findField(clazz, "session");
        ReflectionUtils.makeAccessible(this.field_session);  //type is WebSocketSession

        this.transferFn = (webSocketSessionHolder) -> (WebSocketSession) ReflectionUtils.getField(field_session, webSocketSessionHolder);

    }


    public Map findAll() {
        Function fn = k -> webSocketSessionMap.get(k);
        return webSocketSessionMap.keySet().stream()
                .collect(Collectors.toUnmodifiableMap(Function.identity(), fn.andThen(transferFn)));
    }

    public WebSocketSession getOrNull(String id) {
        return Optional.ofNullable(webSocketSessionMap.getOrDefault(id, null))
                .filter(Objects::nonNull)
                .map(transferFn)
                .orElseGet(() -> null);
    }

}

结合 springboot-actuator 暴露 session

import com.ft.suse.websocket.session.WebSocketSessionContainer;
import lombok.*;
import org.springframework.boot.actuate.endpoint.annotation.Endpoint;
import org.springframework.boot.actuate.endpoint.annotation.ReadOperation;
import org.springframework.boot.actuate.endpoint.annotation.Selector;
import org.springframework.boot.actuate.endpoint.annotation.WriteOperation;
import org.springframework.http.HttpHeaders;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.WebSocketExtension;
import org.springframework.web.socket.WebSocketSession;

import javax.inject.Inject;
import java.io.IOException;
import java.net.InetSocketAddress;
import java.net.URI;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.stream.Collectors;

/**
 * http or jmx
 */
@Endpoint(id = "websocket-session")
@Component
public class WebSocketSessionEndpoint {

    private final WebSocketSessionContainer wsSessionContainer;

    @Inject
    public WebSocketSessionEndpoint(WebSocketSessionContainer wsSessionContainer) {
        this.wsSessionContainer = wsSessionContainer;
    }


    @ReadOperation
    public Map sessions() {
        return wsSessionContainer.findAll().entrySet().stream()
                .filter(Objects::nonNull)
                .collect(Collectors.toMap(e -> e.getKey(), e -> from(e.getValue())));
    }

    @ReadOperation
    public MySessionDescriptor sessionEntry(@Selector String id) {
        return Optional.ofNullable(wsSessionContainer.getOrNull(id))
                .filter(Objects::nonNull)
                .map(this::from)
                .orElseGet(() -> null);
    }

    @WriteOperation
    public String closeSession(@Selector String id) {
        WebSocketSession wsSession = wsSessionContainer.getOrNull(id);
        if (wsSession != null) {
            if (wsSession.isOpen()) {
                try {
                    wsSession.close();
                    return "[success]";
                } catch (IOException e) {
//                e.printStackTrace();
                    return "[fail]:" + e.getMessage();
                }
            } else {
                return "[fail]:session is already closed";
            }
        }
        return "[fail]:session not existed";
    }

    private MySessionDescriptor from(WebSocketSession session) {
        return MySessionDescriptor.builder()
                .id(session.getId())
                .uri(session.getUri())
                .username(session.getPrincipal().getName())
                .acceptedProtocol(session.getAcceptedProtocol())
                .binaryMessageSizeLimit(session.getBinaryMessageSizeLimit())
                .extensions(session.getExtensions())
                .handshakeHeaders(session.getHandshakeHeaders())
                .localAddress(session.getLocalAddress())
                .remoteAddress(session.getRemoteAddress())
                .open(session.isOpen())
                .textMessageSizeLimit(session.getTextMessageSizeLimit())
                .build();
    }

    @Builder
    @NoArgsConstructor
    @AllArgsConstructor
    @Setter
    @Getter
    public static class MySessionDescriptor {
        private String                   id;
        private URI                      uri;
        private String                   username;
        private boolean                  open;
        private InetSocketAddress        remoteAddress;
        private InetSocketAddress        localAddress;
        private String                   acceptedProtocol;
        private Map      attributes;
        private HttpHeaders              handshakeHeaders;
        private int                      textMessageSizeLimit;
        private int                      binaryMessageSizeLimit;
        private List extensions;
    }
}

import org.springframework.boot.actuate.endpoint.web.annotation.RestControllerEndpoint;
import org.springframework.stereotype.Component;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.PathVariable;

import javax.inject.Inject;
import java.util.Map;

@Component
@RestControllerEndpoint(id = "websocket-session-rest")
public class WebSocketSessionController {
    private final WebSocketSessionEndpoint delegate;

    @Inject
    public WebSocketSessionController(WebSocketSessionEndpoint delegate) {
        this.delegate = delegate;
    }

    @GetMapping("")
    public Map sessions() {
        return delegate.sessions();
    }

    @GetMapping("/{id}")
    public WebSocketSessionEndpoint.MySessionDescriptor sessionEntry(@PathVariable("id") String id) {
        return delegate.sessionEntry(id);
    }

    @GetMapping("/{id}/close")
    public String closeSession(@PathVariable("id") String id) {
        return delegate.closeSession(id);
    }
}