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 new file mode 100644 index 0000000000..6408729877 --- /dev/null +++ b/common/queue/src/main/java/org/thingsboard/server/queue/memory/DefaultInMemoryStorage.java @@ -0,0 +1,75 @@ +/** + * Copyright © 2016-2022 The Thingsboard Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.thingsboard.server.queue.memory; + +import lombok.extern.slf4j.Slf4j; +import org.springframework.stereotype.Component; +import org.thingsboard.server.queue.TbQueueMsg; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.LinkedBlockingQueue; + +@Component +@Slf4j +public final class DefaultInMemoryStorage implements InMemoryStorage { + private final ConcurrentHashMap> storage = new ConcurrentHashMap<>(); + + @Override + public void printStats() { + if (log.isDebugEnabled()) { + storage.forEach((topic, queue) -> { + if (queue.size() > 0) { + log.debug("[{}] Queue Size [{}]", topic, queue.size()); + } + }); + } + } + + @Override + public int getLagTotal() { + return storage.values().stream().map(BlockingQueue::size).reduce(0, Integer::sum); + } + + @Override + public boolean put(String topic, TbQueueMsg msg) { + return storage.computeIfAbsent(topic, (t) -> new LinkedBlockingQueue<>()).add(msg); + } + + @SuppressWarnings("unchecked") + @Override + public List get(String topic) throws InterruptedException { + 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; + } + return Collections.singletonList((T) firstMsg); + } + } + return Collections.emptyList(); + } + +} diff --git a/common/queue/src/main/java/org/thingsboard/server/queue/memory/InMemoryStorage.java b/common/queue/src/main/java/org/thingsboard/server/queue/memory/InMemoryStorage.java index 11b97100ee..0fdfebbfff 100644 --- a/common/queue/src/main/java/org/thingsboard/server/queue/memory/InMemoryStorage.java +++ b/common/queue/src/main/java/org/thingsboard/server/queue/memory/InMemoryStorage.java @@ -15,59 +15,18 @@ */ package org.thingsboard.server.queue.memory; -import lombok.extern.slf4j.Slf4j; -import org.springframework.stereotype.Component; import org.thingsboard.server.queue.TbQueueMsg; -import java.util.ArrayList; -import java.util.Collections; import java.util.List; -import java.util.concurrent.BlockingQueue; -import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.LinkedBlockingQueue; -@Component -@Slf4j -public final class InMemoryStorage { - private final ConcurrentHashMap> storage = new ConcurrentHashMap<>(); +public interface InMemoryStorage { - public void printStats() { - storage.forEach((topic, queue) -> { - if (queue.size() > 0) { - log.debug("[{}] Queue Size [{}]", topic, queue.size()); - } - }); - } + void printStats(); - public int getLagTotal() { - return storage.values().stream().map(BlockingQueue::size).reduce(0, Integer::sum); - } + int getLagTotal(); - public boolean put(String topic, TbQueueMsg msg) { - return storage.computeIfAbsent(topic, (t) -> new LinkedBlockingQueue<>()).add(msg); - } + boolean put(String topic, TbQueueMsg msg); - 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); - } - } else { - entities = Collections.emptyList(); - } - return entities; - } - return Collections.emptyList(); - } + List get(String topic) throws InterruptedException; } 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 new file mode 100644 index 0000000000..b269e55888 --- /dev/null +++ b/common/queue/src/test/java/org/thingsboard/server/queue/memory/DefaultInMemoryStorageTest.java @@ -0,0 +1,115 @@ +/** + * Copyright © 2016-2022 The Thingsboard Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.thingsboard.server.queue.memory; + +import com.google.gson.Gson; +import lombok.extern.slf4j.Slf4j; +import org.junit.Test; +import org.thingsboard.server.queue.TbQueueMsg; +import org.thingsboard.server.queue.common.DefaultTbQueueMsg; + +import java.util.ArrayList; +import java.util.List; +import java.util.UUID; + +import static org.assertj.core.api.Assertions.assertThat; +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(); + + @Test + public void givenStorage_whenGetLagTotal_thenReturnInteger() throws InterruptedException { + assertThat(storage.getLagTotal()).isEqualTo(0); + storage.put("main", mock(TbQueueMsg.class)); + assertThat(storage.getLagTotal()).isEqualTo(1); + storage.put("main", mock(TbQueueMsg.class)); + assertThat(storage.getLagTotal()).isEqualTo(2); + storage.put("hp", mock(TbQueueMsg.class)); + assertThat(storage.getLagTotal()).isEqualTo(3); + storage.get("main"); + assertThat(storage.getLagTotal()).isEqualTo(1); + } + + @Test + 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(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); + } + +} diff --git a/common/queue/src/test/java/org/thingsboard/server/queue/memory/InMemoryStorageTest.java b/common/queue/src/test/java/org/thingsboard/server/queue/memory/InMemoryStorageTest.java deleted file mode 100644 index d8911c3fd9..0000000000 --- a/common/queue/src/test/java/org/thingsboard/server/queue/memory/InMemoryStorageTest.java +++ /dev/null @@ -1,42 +0,0 @@ -/** - * Copyright © 2016-2022 The Thingsboard Authors - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package org.thingsboard.server.queue.memory; - -import org.junit.After; -import org.junit.Before; -import org.junit.Test; -import org.thingsboard.server.queue.TbQueueMsg; - -import static org.assertj.core.api.Assertions.assertThat; -import static org.mockito.Mockito.mock; - -public class InMemoryStorageTest { - - InMemoryStorage storage = new InMemoryStorage(); - - @Test - public void givenStorage_whenGetLagTotal_thenReturnInteger() throws InterruptedException { - assertThat(storage.getLagTotal()).isEqualTo(0); - storage.put("main", mock(TbQueueMsg.class)); - assertThat(storage.getLagTotal()).isEqualTo(1); - storage.put("main", mock(TbQueueMsg.class)); - assertThat(storage.getLagTotal()).isEqualTo(2); - storage.put("hp", mock(TbQueueMsg.class)); - assertThat(storage.getLagTotal()).isEqualTo(3); - storage.get("main"); - assertThat(storage.getLagTotal()).isEqualTo(1); - } -} \ No newline at end of file