From f8a675118292d5e48f21f82251de250a2d5aeeba Mon Sep 17 00:00:00 2001 From: Sergey Matvienko Date: Fri, 15 Apr 2022 16:40:32 +0300 Subject: [PATCH] InMemoryStorage performance improved. Many test cases added since it is essential piece of code. --- .../queue/memory/DefaultInMemoryStorage.java | 28 ++++---- .../memory/DefaultInMemoryStorageTest.java | 66 ++++++++++++++++--- 2 files changed, 70 insertions(+), 24 deletions(-) diff --git a/common/queue/src/main/java/org/thingsboard/server/queue/memory/DefaultInMemoryStorage.java b/common/queue/src/main/java/org/thingsboard/server/queue/memory/DefaultInMemoryStorage.java index f852974dd5..64694a817c 100644 --- a/common/queue/src/main/java/org/thingsboard/server/queue/memory/DefaultInMemoryStorage.java +++ b/common/queue/src/main/java/org/thingsboard/server/queue/memory/DefaultInMemoryStorage.java @@ -67,26 +67,22 @@ public final class DefaultInMemoryStorage implements InMemoryStorage { return storage.computeIfAbsent(topic, (t) -> new LinkedBlockingQueue<>()).add(msg); } + @SuppressWarnings("unchecked") @Override public List get(String topic) throws InterruptedException { - if (storage.containsKey(topic)) { - List entities; - @SuppressWarnings("unchecked") - T first = (T) storage.get(topic).poll(); - if (first != null) { - entities = new ArrayList<>(); - entities.add(first); - List otherList = new ArrayList<>(); - storage.get(topic).drainTo(otherList, 999); - for (TbQueueMsg other : otherList) { - @SuppressWarnings("unchecked") - T entity = (T) other; - entities.add(entity); + final BlockingQueue queue = storage.get(topic); + if (queue != null) { + final TbQueueMsg firstMsg = queue.poll(); + if (firstMsg != null) { + final int queueSize = queue.size(); + if (queueSize > 0) { + final List entities = new ArrayList<>(Math.min(queueSize, 999) + 1); + entities.add(firstMsg); + queue.drainTo(entities, 999); + return (List) entities; } - } else { - entities = Collections.emptyList(); + return Collections.singletonList((T) firstMsg); } - return entities; } return Collections.emptyList(); } diff --git a/common/queue/src/test/java/org/thingsboard/server/queue/memory/DefaultInMemoryStorageTest.java b/common/queue/src/test/java/org/thingsboard/server/queue/memory/DefaultInMemoryStorageTest.java index 65d9da1b07..b269e55888 100644 --- a/common/queue/src/test/java/org/thingsboard/server/queue/memory/DefaultInMemoryStorageTest.java +++ b/common/queue/src/test/java/org/thingsboard/server/queue/memory/DefaultInMemoryStorageTest.java @@ -30,6 +30,9 @@ import static org.mockito.Mockito.mock; @Slf4j public class DefaultInMemoryStorageTest { + static final int MAX_POLL_SIZE = 1000; + final Gson gson = new Gson(); + final String topic = "tb_core_notification.tb-node-0"; InMemoryStorage storage = new DefaultInMemoryStorage(); @@ -47,19 +50,66 @@ public class DefaultInMemoryStorageTest { } @Test - public void givenQueue_whenPoll_thenReturnList() throws InterruptedException { - Gson gson = new Gson(); - String topic = "tb_core_notification.tb-node-0"; - List msgs = new ArrayList<>(1001); - for (int i = 0; i < 1001; i++) { + public void givenQueueWithMoreThenBatchSize_whenPoll_thenReturnFullListAndSecondList() throws InterruptedException { + List msgs = new ArrayList<>(MAX_POLL_SIZE + 1); + for (int i = 0; i < MAX_POLL_SIZE + 1; i++) { DefaultTbQueueMsg msg = gson.fromJson("{\"key\": \"" + UUID.randomUUID() + "\"}", DefaultTbQueueMsg.class); msgs.add(msg); storage.put(topic, msg); } - assertThat(storage.getLagTotal()).as("total lag is 1001").isEqualTo(1001); - assertThat(storage.get(topic)).as("poll exactly 1000 msgs").isEqualTo(msgs.subList(0, 1000)); - assertThat(storage.get(topic)).as("poll last 1 message").isEqualTo(msgs.subList(1000, 1001)); + assertThat(storage.getLagTotal()).as("total lag is 1001").isEqualTo(MAX_POLL_SIZE + 1); + assertThat(storage.get(topic)).as("poll exactly 1000 msgs").isEqualTo(msgs.subList(0, MAX_POLL_SIZE)); + assertThat(storage.get(topic)).as("poll last 1 message").isEqualTo(msgs.subList(MAX_POLL_SIZE, MAX_POLL_SIZE + 1)); assertThat(storage.getLagTotal()).as("total lag is zero").isEqualTo(0); } + + private void testPollOnce(final int msgCount) throws InterruptedException { + List msgs = new ArrayList<>(msgCount); + for (int i = 0; i < msgCount; i++) { + DefaultTbQueueMsg msg = gson.fromJson("{\"key\": \"" + UUID.randomUUID() + "\"}", DefaultTbQueueMsg.class); + msgs.add(msg); + storage.put(topic, msg); + } + + assertThat(storage.getLagTotal()).as("total lag before poll").isEqualTo(msgCount); + assertThat(storage.get(topic)).as("polled exactly msgs").isEqualTo(msgs.subList(0, msgCount)); + assertThat(storage.getLagTotal()).as("final lag is zero").isEqualTo(0); + } + + @Test + public void givenQueueWithExactBatchSize_whenPoll_thenReturnExactBatchSizeList() throws InterruptedException { + testPollOnce(MAX_POLL_SIZE); + } + + @Test + public void givenQueueWithExactBatchSizeMinusOne_whenPoll_thenReturnCorrectSizeList() throws InterruptedException { + testPollOnce(MAX_POLL_SIZE - 1); + } + + @Test + public void givenQueueWithExactBatchSizeMinusTen_whenPoll_thenReturnCorrectSizeList() throws InterruptedException { + testPollOnce(MAX_POLL_SIZE - 10); + } + + @Test + public void givenQueueEmpty_whenPoll_thenReturnEmptyList() throws InterruptedException { + testPollOnce(0); + } + + @Test + public void givenQueueWithSingleMessage_whenPoll_thenReturnSingletonList() throws InterruptedException { + testPollOnce(1); + } + + @Test + public void givenQueueWithTwoMessages_whenPoll_thenReturnCorrectSizeList() throws InterruptedException { + testPollOnce(2); + } + + @Test + public void givenQueueWithTenMessages_whenPoll_thenReturnCorrectSizeList() throws InterruptedException { + testPollOnce(10); + } + }