advancedxy commented on code in PR #811:
URL: https://github.com/apache/incubator-uniffle/pull/811#discussion_r1169594377
##########
client-spark/common/src/test/java/org/apache/spark/shuffle/writer/WriteBufferManagerTest.java:
##########
@@ -221,6 +227,58 @@ public void buildBlockEventsTest() {
assertEquals(3, events.size());
}
+ @ParameterizedTest
+ @ValueSource(longs = { 31, 63, 100000 })
+ public void concurrentAddAndSpill(long spillSize) {
+ SparkConf conf = getConf();
+ conf.set("spark.rss.client.send.size.limit", "1000");
+ conf.set("spark.rss.writer.buffer.spill.size", String.valueOf(spillSize));
+ TaskMemoryManager mockTaskMemoryManager = mock(TaskMemoryManager.class);
+ BufferManagerOptions bufferOptions = new BufferManagerOptions(conf);
+
+ AtomicLong processedSize = new AtomicLong();
+ Function<AddBlockEvent, CompletableFuture<Long>> spillFunc = event -> {
+ event.getProcessedCallbackChain().stream().forEach(x -> x.run());
+ long size = event.getShuffleDataInfoList().stream().mapToLong(x ->
x.getFreeMemory()).sum();
+ processedSize.getAndAdd(size);
+ return CompletableFuture.completedFuture(size);
+ };
+
+ WriteBufferManager wbm = new WriteBufferManager(
+ 0, "taskId_spillTest", 0, bufferOptions, new KryoSerializer(conf),
+ Maps.newHashMap(), mockTaskMemoryManager, new ShuffleWriteMetrics(),
+ RssSparkConfig.toRssConf(conf), spillFunc);
+ WriteBufferManager spyManager = spy(wbm);
+ doReturn(512L).when(spyManager).acquireMemory(anyLong());
+
+ AtomicLong expectedSize = new AtomicLong(0L);
+ final String testKey = "Key";
+ final String testValue = "Value";
+ AtomicBoolean produceFinished = new AtomicBoolean(false);
+ AtomicBoolean spillFinished = new AtomicBoolean(false);
+
+ new Thread(() -> {
+ IntStream.range(0, 1000).forEach(x -> {
+ List<ShuffleBlockInfo> blockInfos = spyManager.addRecord(x, testKey,
testValue);
+ blockInfos.stream().forEach(blockInfo ->
processedSize.addAndGet(blockInfo.getFreeMemory()));
+ expectedSize.addAndGet(32);
+ });
+ produceFinished.set(true);
+ }).start();
+
+
+ new Thread(() -> {
+ while (!produceFinished.get()) {
+ spyManager.spill(100000, mock(WriteBufferManager.class));
+ }
+ spyManager.spill(100000, mock(WriteBufferManager.class));
+ spillFinished.set(true);
+ }).start();
Review Comment:
is it possible to start multiple spill threads instead of one?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
