| 1242 | } |
| 1243 | |
| 1244 | void TmpFileMgrTest::TestCompressBufferManagement( |
| 1245 | bool spill_to_remote, bool wait_upload) { |
| 1246 | // Data string should be long and redundant enough to be compressible. |
| 1247 | const string DATA = "the quick brown fox jumped over the lazy dog" |
| 1248 | "the fast red fox leaped over the sleepy dog"; |
| 1249 | const string BIG_DATA = DATA + DATA; |
| 1250 | string data = DATA; |
| 1251 | string big_data = BIG_DATA; |
| 1252 | FLAGS_disk_spill_compression_buffer_limit_bytes = 2 * data.size(); |
| 1253 | // Limit compression buffers to not quite enough for two compression buffers |
| 1254 | // for 'data' (since compression buffers need to be slightly larger than |
| 1255 | // uncompressed data). |
| 1256 | TmpFileMgr tmp_file_mgr; |
| 1257 | if (spill_to_remote) { |
| 1258 | RemoveAndCreateDirs(vector<string>{LOCAL_BUFFER_PATH}); |
| 1259 | ASSERT_OK(tmp_file_mgr.InitCustom(vector<string>{remote_url_, LOCAL_BUFFER_PATH}, |
| 1260 | true, "lz4", true, metrics_.get())); |
| 1261 | } else { |
| 1262 | ASSERT_OK(tmp_file_mgr.InitCustom( |
| 1263 | vector<string>{"/tmp/tmp-file-mgr-test.1"}, true, "lz4", true, metrics_.get())); |
| 1264 | } |
| 1265 | MemTracker* compressed_buffer_tracker = tmp_file_mgr.compressed_buffer_tracker(); |
| 1266 | TmpFileGroup file_group(&tmp_file_mgr, io_mgr(), profile_, TUniqueId()); |
| 1267 | MemRange data_mem_range(reinterpret_cast<uint8_t*>(&data[0]), data.size()); |
| 1268 | MemRange big_data_mem_range(reinterpret_cast<uint8_t*>(&big_data[0]), big_data.size()); |
| 1269 | |
| 1270 | // Start a write in flight, which should encrypt the data and write it to disk. |
| 1271 | unique_ptr<TmpWriteHandle> compressed_handle, uncompressed_handle; |
| 1272 | WriteRange::WriteDoneCallback callback = |
| 1273 | bind(mem_fn(&TmpFileMgrTest::SignalCallback), this, _1); |
| 1274 | ASSERT_OK(file_group.Write(data_mem_range, callback, &compressed_handle)); |
| 1275 | EXPECT_TRUE(compressed_handle->is_compressed()); |
| 1276 | int mem_consumption_after_first_write = compressed_buffer_tracker->peak_consumption(); |
| 1277 | EXPECT_GT(mem_consumption_after_first_write, 0) |
| 1278 | << "Compressed buffer memory should be consumed for in-flight writes"; |
| 1279 | WaitForWrite(compressed_handle.get()); |
| 1280 | EXPECT_EQ(0, compressed_buffer_tracker->consumption()) |
| 1281 | << "No memory should be consumed when no reads or writes in-flight"; |
| 1282 | |
| 1283 | // Spot-check that bytes written counters were updated correctly. |
| 1284 | EXPECT_GT(file_group.bytes_written_counter_->value(), 0); |
| 1285 | EXPECT_GT(file_group.uncompressed_bytes_written_counter_->value(), |
| 1286 | file_group.bytes_written_counter_->value()); |
| 1287 | |
| 1288 | // This data range is larger than the memory limit and falls back to uncompressed |
| 1289 | // writes. |
| 1290 | ASSERT_OK(file_group.Write(big_data_mem_range, callback, &uncompressed_handle)); |
| 1291 | EXPECT_EQ(uncompressed_handle->data_len(), uncompressed_handle->on_disk_len()); |
| 1292 | EXPECT_FALSE(uncompressed_handle->is_compressed()); |
| 1293 | EXPECT_LE(compressed_buffer_tracker->consumption(), mem_consumption_after_first_write) |
| 1294 | << "Second write should have fallen back to uncompressed writes"; |
| 1295 | |
| 1296 | WaitForWrite(uncompressed_handle.get()); |
| 1297 | WaitForCallbacks(2); |
| 1298 | EXPECT_EQ(0, compressed_buffer_tracker->consumption()) |
| 1299 | << "No memory should be consumed when no reads or writes in-flight"; |
| 1300 | |
| 1301 | if (spill_to_remote && wait_upload) { |
nothing calls this directly
no test coverage detected