Handling ‘State’ in Java WebSocket Applications

By and large, there are two kinds of states in a WebSocket application

• User/client specific: related to a connected user/Session e.g. user ID, list of subscriptions, last message received, etc.
• Global: state which is relevant across your application and something which all connected users/Sessions might be able to use.

User Specific State

This can be handled using getUserProperties method on the Session object – this exposes a Map, which you can use to store anything (Object type) using a String type key:

@OnOpen
public void opened(@PathParam("userid") String id, Session peer){
    peer.getUserProperties().put("USER_ID" , id); //it's possible to store the ID as a member variable as well.. this is just an example
}

@OnMessage
public void helloUser(Session peer, String message){
    String id = (String) peer.getUserProperties().get("USER_ID");
    peer.getBasicRemote().sendText("Hello "+ id + "! You sent "+ message);
}

Global State

There are multiple options here as well. Please note that these are scoped to a specific Endpoint

getUserProperties in EndpointConfig – it exposes the same Map interface as the one in Session. Since the WebSocket runtime creates a single instance of an EndpointConfig object per Endpoint , it can be used a global state store:

private EndpointCondig epCfg;
private static List<Session> peers = ...;

@OnOpen
public void open(@PathParam("userid") String id, Session peer, EndpointConfig epCfg){
    this.epCfg = epCfg;
    peers.add(peer);
    this.epCfg.getUserProperties().put(peer.getId(), id); //store mapping of WebSocket Session ID to user ID
}

@OnMessage
public void broadcast(Session from, String msg){
    String senderID = (String) this.epCfg.getUserProperties().get(from.getID()); //check the mapping
    for(Session peer : peers) { //loop over ALL connected clients
        if(peer.isOpen()){
            peer.getBasicRemote().sendText("Message from User "+ senderID + " - " + msg);
        }
    }
}

Another option is to encapsulate some of the common/global logic in a custom Configurator implementation, which can be accessed and used within the endpoint logic:

//custom Configurator implementation which maps (authenticated) user name to a token (sent via HTTP header)

public class TokenStore extends ServerEndpointConfig.Configurator {
    Map<String, String> userTokens;

    public TokenStore() {
        userTokens = new ConcurrentHashMap<>();
    }

    @Override
    public void modifyHandshake(ServerEndpointConfig sec, HandshakeRequest request, HandshakeResponse response) {
        String token = request.getHeaders().get("token").get(0);
        String name = request.getUserPrincipal().getName();
        userTokens.put(name, token);
    }

    public Map<String, String> getUserTokens(){
        return Collections.unmodifiableMap(userTokens);
    }

}

//the WebSocket (server) endpoint implementation which makes use of the token store

@ServerEndpoint(value = "/service/{id}",configurator = TokenStore.class)
public class BroadcastService {
    private String token;

    @OnOpen
    public void test(@PathParam("id") String id, EndpointConfig cfg) { //injeted config by runtime
        ServerEndpoint sCfg = (ServerEndpoint) cfg; //cast
        TokenStore store = sCfg.getConfigurator(); //get custom implementation instance
        token = cfgur.getUserTokens().get(id); //extract token and store as a member variable
    }
}

Further Reading

• Java WebSocket API specification
• Java WebSocket API Handbook
• Other WebSocket blogs

Cheers!