diff --git a/spring-boot-samples/spring-boot-sample-websocket/src/main/java/org/springframework/boot/samples/websocket/snake/SnakeTimer.java b/spring-boot-samples/spring-boot-sample-websocket/src/main/java/org/springframework/boot/samples/websocket/snake/SnakeTimer.java index 0b9003e95c..99cc83f014 100644 --- a/spring-boot-samples/spring-boot-sample-websocket/src/main/java/org/springframework/boot/samples/websocket/snake/SnakeTimer.java +++ b/spring-boot-samples/spring-boot-sample-websocket/src/main/java/org/springframework/boot/samples/websocket/snake/SnakeTimer.java @@ -22,6 +22,7 @@ import java.util.Iterator; import java.util.Timer; import java.util.TimerTask; import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.CopyOnWriteArrayList; import org.apache.juli.logging.Log; import org.apache.juli.logging.LogFactory; @@ -78,14 +79,16 @@ public class SnakeTimer { } public static void broadcast(String message) throws Exception { - for (Snake snake : SnakeTimer.getSnakes()) { - try { - snake.sendMessage(message); - } - catch (Throwable ex) { - // if Snake#sendMessage fails the client should be removed - } - } + Collection snakes = new CopyOnWriteArrayList<>(SnakeTimer.getSnakes()); + for (Snake snake : snakes) { + try { + snake.sendMessage(message); + } + catch (Throwable ex) { + // if Snake#sendMessage fails the client is removed + removeSnake(snake); + } + } } diff --git a/spring-boot-samples/spring-boot-sample-websocket/src/test/java/org/springframework/boot/samples/websocket/snake/SnakeTimerTests.java b/spring-boot-samples/spring-boot-sample-websocket/src/test/java/org/springframework/boot/samples/websocket/snake/SnakeTimerTests.java new file mode 100644 index 0000000000..8910711329 --- /dev/null +++ b/spring-boot-samples/spring-boot-sample-websocket/src/test/java/org/springframework/boot/samples/websocket/snake/SnakeTimerTests.java @@ -0,0 +1,24 @@ +package org.springframework.boot.samples.websocket.snake; + +import org.junit.Test; + +import java.io.IOException; + +import static org.hamcrest.CoreMatchers.is; +import static org.junit.Assert.assertThat; +import static org.mockito.Matchers.anyString; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; + +public class SnakeTimerTests { + + @Test + public void removeDysfunctionalSnakes() throws Exception { + Snake snake = mock(Snake.class); + doThrow(new IOException()).when(snake).sendMessage(anyString()); + SnakeTimer.addSnake(snake); + + SnakeTimer.broadcast(""); + assertThat(SnakeTimer.getSnakes().size(), is(0)); + } +}