1 // Copyright (c) 2011-present, Facebook, Inc.  All rights reserved.
2 //  This source code is licensed under both the GPLv2 (found in the
3 //  COPYING file in the root directory) and Apache 2.0 License
4 //  (found in the LICENSE.Apache file in the root directory).
5 
6 package org.rocksdb;
7 
8 import org.junit.ClassRule;
9 import org.junit.Rule;
10 import org.junit.Test;
11 import org.junit.rules.TemporaryFolder;
12 
13 import java.nio.charset.StandardCharsets;
14 import java.util.*;
15 
16 import static org.assertj.core.api.Assertions.assertThat;
17 
18 public class MemoryUtilTest {
19 
20   private static final String MEMTABLE_SIZE = "rocksdb.size-all-mem-tables";
21   private static final String UNFLUSHED_MEMTABLE_SIZE = "rocksdb.cur-size-all-mem-tables";
22   private static final String TABLE_READERS = "rocksdb.estimate-table-readers-mem";
23 
24   private final byte[] key = "some-key".getBytes(StandardCharsets.UTF_8);
25   private final byte[] value = "some-value".getBytes(StandardCharsets.UTF_8);
26 
27   @ClassRule
28   public static final RocksNativeLibraryResource ROCKS_NATIVE_LIBRARY_RESOURCE =
29       new RocksNativeLibraryResource();
30 
31   @Rule public TemporaryFolder dbFolder1 = new TemporaryFolder();
32   @Rule public TemporaryFolder dbFolder2 = new TemporaryFolder();
33 
34   /**
35    * Test MemoryUtil.getApproximateMemoryUsageByType before and after a put + get
36    */
37   @Test
getApproximateMemoryUsageByType()38   public void getApproximateMemoryUsageByType() throws RocksDBException {
39     try (final Cache cache = new LRUCache(8 * 1024 * 1024);
40          final Options options =
41                  new Options()
42                          .setCreateIfMissing(true)
43                          .setTableFormatConfig(new BlockBasedTableConfig().setBlockCache(cache));
44          final FlushOptions flushOptions =
45                  new FlushOptions().setWaitForFlush(true);
46          final RocksDB db =
47                  RocksDB.open(options, dbFolder1.getRoot().getAbsolutePath())) {
48 
49       List<RocksDB> dbs = new ArrayList<>(1);
50       dbs.add(db);
51       Set<Cache> caches = new HashSet<>(1);
52       caches.add(cache);
53       Map<MemoryUsageType, Long> usage = MemoryUtil.getApproximateMemoryUsageByType(dbs, caches);
54 
55       assertThat(usage.get(MemoryUsageType.kMemTableTotal)).isEqualTo(
56               db.getAggregatedLongProperty(MEMTABLE_SIZE));
57       assertThat(usage.get(MemoryUsageType.kMemTableUnFlushed)).isEqualTo(
58               db.getAggregatedLongProperty(UNFLUSHED_MEMTABLE_SIZE));
59       assertThat(usage.get(MemoryUsageType.kTableReadersTotal)).isEqualTo(
60               db.getAggregatedLongProperty(TABLE_READERS));
61       assertThat(usage.get(MemoryUsageType.kCacheTotal)).isEqualTo(0);
62 
63       db.put(key, value);
64       db.flush(flushOptions);
65       db.get(key);
66 
67       usage = MemoryUtil.getApproximateMemoryUsageByType(dbs, caches);
68       assertThat(usage.get(MemoryUsageType.kMemTableTotal)).isGreaterThan(0);
69       assertThat(usage.get(MemoryUsageType.kMemTableTotal)).isEqualTo(
70               db.getAggregatedLongProperty(MEMTABLE_SIZE));
71       assertThat(usage.get(MemoryUsageType.kMemTableUnFlushed)).isGreaterThan(0);
72       assertThat(usage.get(MemoryUsageType.kMemTableUnFlushed)).isEqualTo(
73               db.getAggregatedLongProperty(UNFLUSHED_MEMTABLE_SIZE));
74       assertThat(usage.get(MemoryUsageType.kTableReadersTotal)).isGreaterThan(0);
75       assertThat(usage.get(MemoryUsageType.kTableReadersTotal)).isEqualTo(
76               db.getAggregatedLongProperty(TABLE_READERS));
77       assertThat(usage.get(MemoryUsageType.kCacheTotal)).isGreaterThan(0);
78 
79     }
80   }
81 
82   /**
83    * Test MemoryUtil.getApproximateMemoryUsageByType with null inputs
84    */
85   @Test
getApproximateMemoryUsageByTypeNulls()86   public void getApproximateMemoryUsageByTypeNulls() throws RocksDBException {
87     Map<MemoryUsageType, Long> usage = MemoryUtil.getApproximateMemoryUsageByType(null, null);
88 
89     assertThat(usage.get(MemoryUsageType.kMemTableTotal)).isEqualTo(null);
90     assertThat(usage.get(MemoryUsageType.kMemTableUnFlushed)).isEqualTo(null);
91     assertThat(usage.get(MemoryUsageType.kTableReadersTotal)).isEqualTo(null);
92     assertThat(usage.get(MemoryUsageType.kCacheTotal)).isEqualTo(null);
93   }
94 
95   /**
96    * Test MemoryUtil.getApproximateMemoryUsageByType with two DBs and two caches
97    */
98   @Test
getApproximateMemoryUsageByTypeMultiple()99   public void getApproximateMemoryUsageByTypeMultiple() throws RocksDBException {
100     try (final Cache cache1 = new LRUCache(1 * 1024 * 1024);
101          final Options options1 =
102                  new Options()
103                          .setCreateIfMissing(true)
104                          .setTableFormatConfig(new BlockBasedTableConfig().setBlockCache(cache1));
105          final RocksDB db1 =
106                  RocksDB.open(options1, dbFolder1.getRoot().getAbsolutePath());
107          final Cache cache2 = new LRUCache(1 * 1024 * 1024);
108          final Options options2 =
109                  new Options()
110                          .setCreateIfMissing(true)
111                          .setTableFormatConfig(new BlockBasedTableConfig().setBlockCache(cache2));
112          final RocksDB db2 =
113                  RocksDB.open(options2, dbFolder2.getRoot().getAbsolutePath());
114          final FlushOptions flushOptions =
115                  new FlushOptions().setWaitForFlush(true);
116 
117     ) {
118       List<RocksDB> dbs = new ArrayList<>(1);
119       dbs.add(db1);
120       dbs.add(db2);
121       Set<Cache> caches = new HashSet<>(1);
122       caches.add(cache1);
123       caches.add(cache2);
124 
125       for (RocksDB db: dbs) {
126         db.put(key, value);
127         db.flush(flushOptions);
128         db.get(key);
129       }
130 
131       Map<MemoryUsageType, Long> usage = MemoryUtil.getApproximateMemoryUsageByType(dbs, caches);
132       assertThat(usage.get(MemoryUsageType.kMemTableTotal)).isEqualTo(
133               db1.getAggregatedLongProperty(MEMTABLE_SIZE) + db2.getAggregatedLongProperty(MEMTABLE_SIZE));
134       assertThat(usage.get(MemoryUsageType.kMemTableUnFlushed)).isEqualTo(
135               db1.getAggregatedLongProperty(UNFLUSHED_MEMTABLE_SIZE) + db2.getAggregatedLongProperty(UNFLUSHED_MEMTABLE_SIZE));
136       assertThat(usage.get(MemoryUsageType.kTableReadersTotal)).isEqualTo(
137               db1.getAggregatedLongProperty(TABLE_READERS) + db2.getAggregatedLongProperty(TABLE_READERS));
138       assertThat(usage.get(MemoryUsageType.kCacheTotal)).isGreaterThan(0);
139 
140     }
141   }
142 
143 }
144