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