Merge pull request #6422 from smatvienko-tb/in-memory-storage-improvements-develop34

[3.4] In memory storage improvements

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<String, BlockingQueue<TbQueueMsg>> 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 <T extends TbQueueMsg> List<T> get(String topic) throws InterruptedException {
+        final BlockingQueue<TbQueueMsg> 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<TbQueueMsg> entities = new ArrayList<>(Math.min(queueSize, 999) + 1);
+                    entities.add(firstMsg);
+                    queue.drainTo(entities, 999);
+                    return (List<T>) entities;
+                }
+                return Collections.singletonList((T) firstMsg);
+            }
+        }
+        return Collections.emptyList();
+    }
+
+}

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
+public interface InMemoryStorage {
-@Slf4j
-public final class InMemoryStorage {
-    private final ConcurrentHashMap<String, BlockingQueue<TbQueueMsg>> storage = new ConcurrentHashMap<>();
 
-    public void printStats() {
+    void printStats();
-        storage.forEach((topic, queue) -> {
-            if (queue.size() > 0) {
-                log.debug("[{}] Queue Size [{}]", topic, queue.size());
-            }
-        });
-    }
 
-    public int getLagTotal() {
+    int getLagTotal();
-        return storage.values().stream().map(BlockingQueue::size).reduce(0, Integer::sum);
-    }
 
-    public boolean put(String topic, TbQueueMsg msg) {
+    boolean put(String topic, TbQueueMsg msg);
-        return storage.computeIfAbsent(topic, (t) -> new LinkedBlockingQueue<>()).add(msg);
-    }
 
-    public <T extends TbQueueMsg> List<T> get(String topic) throws InterruptedException {
+    <T extends TbQueueMsg> List<T> get(String topic) throws InterruptedException;
-        if (storage.containsKey(topic)) {
-            List<T> entities;
-            @SuppressWarnings("unchecked")
-            T first = (T) storage.get(topic).poll();
-            if (first != null) {
-                entities = new ArrayList<>();
-                entities.add(first);
-                List<TbQueueMsg> 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();
-    }
 
 }

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<TbQueueMsg> 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<TbQueueMsg> 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);
+    }
+
+}

common/queue/src/test/java/org/thingsboard/server/queue/memory/InMemoryStorageTest.java

@@ -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);
-    }
-}