added Caffeine cache for Cassandra ts partitions saving

application/src/main/resources/thingsboard.yml

@@ -192,8 +192,9 @@ cassandra:
     read_consistency_level: "${CASSANDRA_READ_CONSISTENCY_LEVEL:ONE}"
     write_consistency_level: "${CASSANDRA_WRITE_CONSISTENCY_LEVEL:ONE}"
     default_fetch_size: "${CASSANDRA_DEFAULT_FETCH_SIZE:2000}"
-    # Specify partitioning size for timestamp key-value storage. Example: MINUTES, HOURS, DAYS, MONTHS,INDEFINITE
+    # Specify partitioning size for timestamp key-value storage. Example: MINUTES, HOURS, DAYS, MONTHS, INDEFINITE
     ts_key_value_partitioning: "${TS_KV_PARTITIONING:MONTHS}"
+    ts_key_value_partitions_max_cache_size: "${TS_KV_PARTITIONS_MAX_CACHE_SIZE:100000}"
     ts_key_value_ttl: "${TS_KV_TTL:0}"
     events_ttl: "${TS_EVENTS_TTL:0}"
     # Specify TTL of debug log in seconds. The current value corresponds to one week

dao/src/main/java/org/thingsboard/server/dao/timeseries/CassandraBaseTimeseriesDao.java

@@ -79,12 +79,17 @@ public class CassandraBaseTimeseriesDao extends AbstractCassandraBaseTimeseriesD
 
     protected static List<Long> FIXED_PARTITION = Arrays.asList(new Long[]{0L});
 
+    private CassandraTsPartitionsCache cassandraTsPartitionsCache;
+
     @Autowired
     private Environment environment;
 
     @Value("${cassandra.query.ts_key_value_partitioning}")
     private String partitioning;
 
+    @Value("${cassandra.query.ts_key_value_partitions_max_cache_size:100000}")
+    private long partitionsCacheSize;
+
     @Value("${cassandra.query.ts_key_value_ttl}")
     private long systemTtl;
 
@@ -111,13 +116,16 @@ public class CassandraBaseTimeseriesDao extends AbstractCassandraBaseTimeseriesD
         super.startExecutor();
         if (!isInstall()) {
             getFetchStmt(Aggregation.NONE, DESC_ORDER);
-        }
-        Optional<NoSqlTsPartitionDate> partition = NoSqlTsPartitionDate.parse(partitioning);
-        if (partition.isPresent()) {
-            tsFormat = partition.get();
-        } else {
-            log.warn("Incorrect configuration of partitioning {}", partitioning);
-            throw new RuntimeException("Failed to parse partitioning property: " + partitioning + "!");
+            Optional<NoSqlTsPartitionDate> partition = NoSqlTsPartitionDate.parse(partitioning);
+            if (partition.isPresent()) {
+                tsFormat = partition.get();
+                if (!isFixedPartitioning() && partitionsCacheSize > 0) {
+                    cassandraTsPartitionsCache = new CassandraTsPartitionsCache(partitionsCacheSize);
+                }
+            } else {
+                log.warn("Incorrect configuration of partitioning {}", partitioning);
+                throw new RuntimeException("Failed to parse partitioning property: " + partitioning + "!");
+            }
         }
     }
 
@@ -168,26 +176,6 @@ public class CassandraBaseTimeseriesDao extends AbstractCassandraBaseTimeseriesD
         return Futures.transform(Futures.allAsList(futures), result -> dataPointDays, MoreExecutors.directExecutor());
     }
 
-    @Override
-    public ListenableFuture<Integer> savePartition(TenantId tenantId, EntityId entityId, long tsKvEntryTs, String key, long ttl) {
-        if (isFixedPartitioning()) {
-            return Futures.immediateFuture(null);
-        }
-        ttl = computeTtl(ttl);
-        long partition = toPartitionTs(tsKvEntryTs);
-        log.debug("Saving partition {} for the entity [{}-{}] and key {}", partition, entityId.getEntityType(), entityId.getId(), key);
-        BoundStatementBuilder stmtBuilder = new BoundStatementBuilder((ttl == 0 ? getPartitionInsertStmt() : getPartitionInsertTtlStmt()).bind());
-        stmtBuilder.setString(0, entityId.getEntityType().name())
-                .setUuid(1, entityId.getId())
-                .setLong(2, partition)
-                .setString(3, key);
-        if (ttl > 0) {
-            stmtBuilder.setInt(4, (int) ttl);
-        }
-        BoundStatement stmt = stmtBuilder.build();
-        return getFuture(executeAsyncWrite(tenantId, stmt), rs -> 0);
-    }
-
     @Override
     public ListenableFuture<Void> remove(TenantId tenantId, EntityId entityId, DeleteTsKvQuery query) {
         long minPartition = toPartitionTs(query.getStartTs());
@@ -461,6 +449,68 @@ public class CassandraBaseTimeseriesDao extends AbstractCassandraBaseTimeseriesD
         return getFuture(executeAsyncWrite(tenantId, stmt), rs -> null);
     }
 
+    @Override
+    public ListenableFuture<Integer> savePartition(TenantId tenantId, EntityId entityId, long tsKvEntryTs, String key, long ttl) {
+        if (isFixedPartitioning()) {
+            return Futures.immediateFuture(null);
+        }
+        ttl = computeTtl(ttl);
+        long partition = toPartitionTs(tsKvEntryTs);
+        if (cassandraTsPartitionsCache == null) {
+            return doSavePartition(tenantId, entityId, key, ttl, partition);
+        } else {
+            CassandraPartitionCacheKey partitionSearchKey = new CassandraPartitionCacheKey(entityId, key, partition);
+            if (!cassandraTsPartitionsCache.has(partitionSearchKey)) {
+                ListenableFuture<Integer> result = doSavePartition(tenantId, entityId, key, ttl, partition);
+                Futures.addCallback(result, new CacheCallback<>(partitionSearchKey), MoreExecutors.directExecutor());
+                return result;
+            } else {
+                return Futures.immediateFuture(0);
+            }
+        }
+    }
+
+    private ListenableFuture<Integer> doSavePartition(TenantId tenantId, EntityId entityId, String key, long ttl, long partition) {
+        log.debug("Saving partition {} for the entity [{}-{}] and key {}", partition, entityId.getEntityType(), entityId.getId(), key);
+        PreparedStatement preparedStatement = ttl == 0 ? getPartitionInsertStmt() : getPartitionInsertTtlStmt();
+        BoundStatement stmt = preparedStatement.bind();
+        stmt.setString(0, entityId.getEntityType().name());
+        stmt.setUuid(1, entityId.getId());
+        stmt.setLong(2, partition);
+        stmt.setString(3, key);
+        if (ttl > 0) {
+            stmt.setInt(4, (int) ttl);
+        }
+//        BoundStatementBuilder stmtBuilder = new BoundStatementBuilder(bind);
+//        stmtBuilder.setString(0, entityId.getEntityType().name())
+//                .setUuid(1, entityId.getId())
+//                .setLong(2, partition)
+//                .setString(3, key);
+//        if (ttl > 0) {
+//            stmtBuilder.setInt(4, (int) ttl);
+//        }
+//        BoundStatement stmt = stmtBuilder.build();
+        return getFuture(executeAsyncWrite(tenantId, stmt), rs -> 0);
+    }
+
+    private class CacheCallback<Void> implements FutureCallback<Void> {
+        private final CassandraPartitionCacheKey key;
+
+        private CacheCallback(CassandraPartitionCacheKey key) {
+            this.key = key;
+        }
+
+        @Override
+        public void onSuccess(Void result) {
+            cassandraTsPartitionsCache.put(key);
+        }
+
+        @Override
+        public void onFailure(Throwable t) {
+
+        }
+    }
+
     private long computeTtl(long ttl) {
         if (systemTtl > 0) {
             if (ttl == 0) {

dao/src/main/java/org/thingsboard/server/dao/timeseries/CassandraPartitionCacheKey.java

@@ -0,0 +1,30 @@
+/**
+ * Copyright © 2016-2020 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.dao.timeseries;
+
+import lombok.AllArgsConstructor;
+import lombok.Data;
+import org.thingsboard.server.common.data.id.EntityId;
+
+@Data
+@AllArgsConstructor
+public class CassandraPartitionCacheKey {
+
+    private EntityId entityId;
+    private String key;
+    private long partition;
+
+}

dao/src/main/java/org/thingsboard/server/dao/timeseries/CassandraTsPartitionsCache.java

@@ -0,0 +1,42 @@
+/**
+ * Copyright © 2016-2020 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.dao.timeseries;
+
+import com.github.benmanes.caffeine.cache.AsyncLoadingCache;
+import com.github.benmanes.caffeine.cache.Caffeine;
+
+import java.util.concurrent.CompletableFuture;
+
+public class CassandraTsPartitionsCache {
+
+    private AsyncLoadingCache<CassandraPartitionCacheKey, Boolean> partitionsCache;
+
+    public CassandraTsPartitionsCache(long maxCacheSize) {
+        this.partitionsCache = Caffeine.newBuilder()
+                .maximumSize(maxCacheSize)
+                .buildAsync(key -> {
+                    throw new IllegalStateException("'get' methods calls are not supported!");
+                });
+    }
+
+    public boolean has(CassandraPartitionCacheKey key) {
+        return partitionsCache.getIfPresent(key) != null;
+    }
+
+    public void put(CassandraPartitionCacheKey key) {
+        partitionsCache.put(key, CompletableFuture.completedFuture(true));
+    }
+}

dao/src/test/java/org/thingsboard/server/dao/nosql/CassandraPartitionsCacheTest.java

@@ -0,0 +1,106 @@
+/**
+ * Copyright © 2016-2020 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.dao.nosql;
+
+import com.datastax.oss.driver.api.core.ConsistencyLevel;
+import com.datastax.oss.driver.api.core.cql.BoundStatement;
+import com.datastax.oss.driver.api.core.cql.PreparedStatement;
+import com.datastax.oss.driver.api.core.cql.Statement;
+import com.google.common.util.concurrent.Futures;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.runners.MockitoJUnitRunner;
+import org.springframework.core.env.Environment;
+import org.springframework.test.util.ReflectionTestUtils;
+import org.thingsboard.server.common.data.id.TenantId;
+import org.thingsboard.server.dao.cassandra.CassandraCluster;
+import org.thingsboard.server.dao.cassandra.guava.GuavaSession;
+import org.thingsboard.server.dao.timeseries.CassandraBaseTimeseriesDao;
+
+import java.util.UUID;
+
+import static org.mockito.Matchers.any;
+import static org.mockito.Matchers.anyString;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+@RunWith(MockitoJUnitRunner.class)
+public class CassandraPartitionsCacheTest {
+
+    private CassandraBaseTimeseriesDao cassandraBaseTimeseriesDao;
+
+    @Mock
+    private Environment environment;
+
+    @Mock
+    private CassandraBufferedRateExecutor rateLimiter;
+
+    @Mock
+    private CassandraCluster cluster;
+
+    @Mock
+    private GuavaSession session;
+
+    @Mock
+    private PreparedStatement preparedStatement;
+
+    @Mock
+    private BoundStatement boundStatement;
+
+    @Before
+    public void setUp() throws Exception {
+        when(cluster.getDefaultReadConsistencyLevel()).thenReturn(ConsistencyLevel.ONE);
+        when(cluster.getDefaultWriteConsistencyLevel()).thenReturn(ConsistencyLevel.ONE);
+        when(cluster.getSession()).thenReturn(session);
+        when(session.prepare(anyString())).thenReturn(preparedStatement);
+        when(preparedStatement.bind()).thenReturn(boundStatement);
+
+        cassandraBaseTimeseriesDao = spy(new CassandraBaseTimeseriesDao());
+
+        ReflectionTestUtils.setField(cassandraBaseTimeseriesDao, "partitioning", "MONTHS");
+        ReflectionTestUtils.setField(cassandraBaseTimeseriesDao, "partitionsCacheSize", 100000);
+        ReflectionTestUtils.setField(cassandraBaseTimeseriesDao, "systemTtl", 0);
+        ReflectionTestUtils.setField(cassandraBaseTimeseriesDao, "setNullValuesEnabled", false);
+        ReflectionTestUtils.setField(cassandraBaseTimeseriesDao, "environment", environment);
+        ReflectionTestUtils.setField(cassandraBaseTimeseriesDao, "rateLimiter", rateLimiter);
+        ReflectionTestUtils.setField(cassandraBaseTimeseriesDao, "cluster", cluster);
+
+        doReturn(Futures.immediateFuture(null)).when(cassandraBaseTimeseriesDao).getFuture(any(TbResultSetFuture.class), any());
+    }
+
+    @Test
+    public void testPartitionSave() throws Exception {
+        cassandraBaseTimeseriesDao.init();
+
+        UUID id = UUID.randomUUID();
+        TenantId tenantId = new TenantId(id);
+        long tsKvEntryTs = System.currentTimeMillis();
+
+        for (int i = 0; i < 50000; i++) {
+            cassandraBaseTimeseriesDao.savePartition(tenantId, tenantId, tsKvEntryTs, "test" + i, 0);
+        }
+
+        for (int i = 0; i < 60000; i++) {
+            cassandraBaseTimeseriesDao.savePartition(tenantId, tenantId, tsKvEntryTs, "test" + i, 0);
+        }
+        verify(cassandraBaseTimeseriesDao, times(60000)).executeAsyncWrite(any(TenantId.class), any(Statement.class));
+    }
+}

dao/src/test/resources/cassandra-test.properties

@@ -54,6 +54,8 @@ cassandra.query.default_fetch_size=2000
 
 cassandra.query.ts_key_value_partitioning=HOURS
 
+cassandra.query.ts_key_value_partitions_max_cache_size=100000
+
 cassandra.query.ts_key_value_ttl=0
 
 cassandra.query.debug_events_ttl=604800