Refactor Client and SseEventBus to use AtomicInteger and AtomicReference for thread-safe operations

Author: ralscha

src/main/java/ch/rasc/sse/eventbus/Client.java

@@ -25,7 +25,7 @@ public class Client {
 
 	private volatile long lastTransfer;
 
-	private final boolean completeAfterMessage;
+	private volatile boolean completeAfterMessage;
 
 	Client(String id, SseEmitter sseEmitter, boolean completeAfterMessage) {
 		this.id = id;
@@ -54,6 +54,10 @@ void updateEmitter(SseEmitter emitter) {
 		this.sseEmitter = emitter;
 	}
 
+	void updateCompleteAfterMessage(boolean completeAfterMessage) {
+		this.completeAfterMessage = completeAfterMessage;
+	}
+
 	boolean isCompleteAfterMessage() {
 		return this.completeAfterMessage;
 	}

src/main/java/ch/rasc/sse/eventbus/ClientEvent.java

@@ -16,6 +16,7 @@
 package ch.rasc.sse.eventbus;
 
 import java.time.Duration;
+import java.util.concurrent.atomic.AtomicInteger;
 
 import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
 import org.springframework.web.servlet.mvc.method.annotation.SseEmitter.SseEventBuilder;
@@ -28,13 +29,13 @@ public class ClientEvent {
 
 	private final String convertedValue;
 
-	private int errorCounter;
+	private final AtomicInteger errorCounter;
 
 	public ClientEvent(Client client, SseEvent event, String convertedValue) {
 		this.client = client;
 		this.event = event;
 		this.convertedValue = convertedValue;
-		this.errorCounter = 0;
+		this.errorCounter = new AtomicInteger(0);
 	}
 
 	public Client getClient() {
@@ -78,11 +79,11 @@ private static void addStringData(SseEventBuilder sseBuilder, String value) {
 	}
 
 	void incErrorCounter() {
-		this.errorCounter++;
+		this.errorCounter.incrementAndGet();
 	}
 
 	public int getErrorCounter() {
-		return this.errorCounter;
+		return this.errorCounter.get();
 	}
 
 }

src/main/java/ch/rasc/sse/eventbus/SseEventBus.java

@@ -27,6 +27,7 @@
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
 
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
@@ -235,25 +236,40 @@ public void registerClient(String clientId, SseEmitter emitter) {
 	 * first message
 	 */
 	public void registerClient(String clientId, SseEmitter emitter, boolean completeAfterMessage) {
+		AtomicReference<SseEmitter> oldEmitter = new AtomicReference<>();
 		this.clients.compute(clientId, (id, existing) -> {
 			if (existing == null) {
 				return new Client(id, emitter, completeAfterMessage);
 			}
+			oldEmitter.set(existing.sseEmitter());
 			existing.updateEmitter(emitter);
+			existing.updateCompleteAfterMessage(completeAfterMessage);
 			return existing;
 		});
+		if (oldEmitter.get() != null) {
+			try {
+				oldEmitter.get().complete();
+			}
+			catch (Exception e) {
+				logger.debug("Error completing old emitter for client " + clientId, e);
+			}
+		}
 	}
 
 	/**
 	 * Unregisters a client and unsubscribes the client from all events.
 	 * @param clientId unique client identifier
 	 */
 	public void unregisterClient(String clientId) {
-		unsubscribeFromAllEvents(clientId);
-		Client removed = this.clients.remove(clientId);
-		if (removed != null) {
+		AtomicReference<SseEmitter> removedEmitter = new AtomicReference<>();
+		this.clients.computeIfPresent(clientId, (id, client) -> {
+			removedEmitter.set(client.sseEmitter());
+			unsubscribeFromAllEvents(id);
+			return null;
+		});
+		if (removedEmitter.get() != null) {
 			try {
-				removed.sseEmitter().complete();
+				removedEmitter.get().complete();
 			}
 			catch (Exception e) {
 				logger.debug("Error completing emitter for client " + clientId, e);
@@ -284,8 +300,8 @@ public void subscribe(String clientId, String event) {
 	 * @param event the event name
 	 */
 	public void subscribeOnly(String clientId, String event) {
-		this.subscriptionRegistry.subscribe(clientId, event);
 		this.unsubscribeFromAllEvents(clientId, event);
+		this.subscriptionRegistry.subscribe(clientId, event);
 	}
 
 	/**
@@ -493,9 +509,30 @@ private void cleanUpClients() {
 					staleClients.add(entry.getKey());
 				}
 			}
-			staleClients.forEach(this::unregisterClient);
-			if (!staleClients.isEmpty()) {
-				this.listener.afterClientsUnregistered(staleClients);
+			Set<String> actuallyRemoved = new HashSet<>();
+			for (String clientId : staleClients) {
+				long recheckExpiration = System.currentTimeMillis() - this.clientExpiration.toMillis();
+				AtomicReference<SseEmitter> removedEmitter = new AtomicReference<>();
+				this.clients.computeIfPresent(clientId, (id, client) -> {
+					if (client.lastTransfer() < recheckExpiration) {
+						removedEmitter.set(client.sseEmitter());
+						unsubscribeFromAllEvents(id);
+						return null;
+					}
+					return client;
+				});
+				if (removedEmitter.get() != null) {
+					actuallyRemoved.add(clientId);
+					try {
+						removedEmitter.get().complete();
+					}
+					catch (Exception e) {
+						logger.debug("Error completing emitter for client " + clientId, e);
+					}
+				}
+			}
+			if (!actuallyRemoved.isEmpty()) {
+				this.listener.afterClientsUnregistered(actuallyRemoved);
 			}
 		}
 	}