content: /*
   * Licensed to the Apache Software Foundation (ASF) under one or more
   * contributor license agreements.  See the NOTICE file distributed with
   * this work for additional information regarding copyright ownership.
   * The ASF licenses this file to You 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.apache.spark.shuffle.celeborn;
  
  import static org.junit.Assert.assertEquals;
  
  import java.io.File;
  import java.io.IOException;
  import java.util.UUID;
  import java.util.concurrent.atomic.AtomicInteger;
  import java.util.concurrent.atomic.LongAdder;
  
  import scala.Product2;
  import scala.Tuple2;
  import scala.collection.Iterator;
  import scala.collection.mutable.ListBuffer;
  
  import org.apache.spark.Partitioner;
  import org.apache.spark.SparkConf;
  import org.apache.spark.SparkVersionUtil;
  import org.apache.spark.TaskContext;
  import org.apache.spark.memory.TaskMemoryManager;
  import org.apache.spark.memory.UnifiedMemoryManager;
  import org.apache.spark.shuffle.ShuffleWriteMetricsReporter;
  import org.apache.spark.shuffle.ShuffleWriter;
  import org.apache.spark.sql.catalyst.InternalRow;
  import org.apache.spark.sql.catalyst.expressions.UnsafeProjection;
  import org.apache.spark.sql.catalyst.expressions.UnsafeRow;
  import org.apache.spark.sql.execution.UnsafeRowSerializer;
  import org.apache.spark.sql.types.DataType;
  import org.apache.spark.sql.types.StringType$;
  import org.apache.spark.unsafe.types.UTF8String;
  import org.junit.Test;
  import org.mockito.Mockito;
  
  import org.apache.celeborn.client.DummyShuffleClient;
  import org.apache.celeborn.client.ShuffleClient;
  import org.apache.celeborn.common.CelebornConf;
  import org.apache.celeborn.common.util.Utils;
  import org.apache.celeborn.reflect.DynConstructors;
  
  public class SortBasedShuffleWriterSuiteJ extends CelebornShuffleWriterSuiteBase {
    @Override
    protected ShuffleWriter<Integer, String> createShuffleWriter(
        CelebornShuffleHandle handle,
        TaskContext context,
        CelebornConf conf,
        ShuffleClient client,
        ShuffleWriteMetricsReporter metrics)
        throws IOException {
      return new SortBasedShuffleWriter<Integer, String, String>(
          handle, context, conf, client, metrics, SendBufferPool.get(4, 30, 60));
    }
  
    private SortBasedShuffleWriter<Integer, String, String> createShuffleWriterWithPusher(
        CelebornShuffleHandle handle,
        TaskContext context,
        CelebornConf conf,
        ShuffleClient client,
        ShuffleWriteMetricsReporter metrics,
        SortBasedPusher pusher)
        throws Exception {
      return new SortBasedShuffleWriter<Integer, String, String>(
          handle, context, conf, client, metrics, SendBufferPool.get(4, 30, 60), pusher);
    }
  
    private SortBasedPusher createSortBasedPusher(
        CelebornConf conf,
        File tempFile,
        int numPartitions,
        ShuffleWriteMetricsReporter metricsReporter)
        throws Exception {
      SparkConf sparkConf = new SparkConf(false).set("spark.buffer.pageSize", "32k");
      UnifiedMemoryManager unifiedMemoryManager = UnifiedMemoryManager.apply(sparkConf, 1);
      TaskMemoryManager taskMemoryManager = new TaskMemoryManager(unifiedMemoryManager, 0);
  
      final ShuffleClient client = new DummyShuffleClient(conf, tempFile);
      conf.set(CelebornConf.CLIENT_PUSH_SORT_USE_ADAPTIVE_MEMORY_THRESHOLD().key(), "true");
      LongAdder[] mapStatusLengths = new LongAdder[numPartitions];
      for (int i = 0; i < numPartitions; i++) {
        mapStatusLengths[i] = new LongAdder();
      }
      SortBasedPusher pusher =
          new SortBasedPusher(
              taskMemoryManager,
              /*shuffleClient=*/ client,
              /*taskContext=*/ taskContext,
              /*shuffleId=*/ 0,
              /*mapId=*/ 0,
              /*attemptNumber=*/ 0,
              /*taskAttemptId=*/ 0,
              /*numMappers=*/ 0,
              /*numPartitions=*/ numPartitions,
              conf,
              metricsReporter::incBytesWritten,
              mapStatusLengths,
              /*pushSortMemoryThreshold=*/ Utils.byteStringAsBytes("32K"),
              SendBufferPool.get(4, 30, 60));
      return pusher;
    }
  
    private String buildRecord(int size) {
      char[] record = new char[size];
      for (int i = 0; i < size; i++) {
        record[i] = 'a';
      }
      return new String(record);
    }
  
    private Iterator<Product2<Integer, UnsafeRow>> getUnsafeRowIterator(
        final int size, int recordSize, final AtomicInteger total, int numPartitions) {
      int current = 0;
      ListBuffer<Product2<Integer, UnsafeRow>> list = new ListBuffer<>();
      while (current < size) {
        int key = total.getAndIncrement();
        String value = buildRecord(recordSize);
        current += value.length();
        ListBuffer<Object> values = new ListBuffer<>();
        values.$plus$eq(UTF8String.fromString(value));
  
        InternalRow row = InternalRow.apply(values.toSeq());
        DataType[] types = new DataType[1];
        types[0] = StringType$.MODULE$;
        UnsafeRow unsafeRow = UnsafeProjection.create(types).apply(row);
  
        list.$plus$eq(new Tuple2<>(key % numPartitions, unsafeRow));
      }
      return list.toIterator();
    }
  
    @Test
    public void testAdaptiveMemoryThreshold() throws Exception {
      int numPartitions = 8;
      final UnsafeRowSerializer serializer = new UnsafeRowSerializer(1, null);
      final CelebornConf conf =
          new CelebornConf()
              .set(CelebornConf.CLIENT_PUSH_SORT_USE_ADAPTIVE_MEMORY_THRESHOLD().key(), "true")
              .set(CelebornConf.CLIENT_PUSH_BUFFER_MAX_SIZE().key(), "32k");
  
      String PartitionIdPassthroughClazz;
      if (SparkVersionUtil.isGreaterThan(3, 3)) {
        PartitionIdPassthroughClazz = "org.apache.spark.PartitionIdPassthrough";
      } else {
        PartitionIdPassthroughClazz = "org.apache.spark.sql.execution.PartitionIdPassthrough";
      }
      DynConstructors.Ctor<Partitioner> partitionIdPassthroughCtor =
          DynConstructors.builder().impl(PartitionIdPassthroughClazz, int.class).build();
      final Partitioner partitioner = partitionIdPassthroughCtor.newInstance(numPartitions);
      Mockito.doReturn(partitioner).when(dependency).partitioner();
      Mockito.doReturn(serializer).when(dependency).serializer();
  
      final File tempFile = new File(tempDir, UUID.randomUUID().toString());
      final CelebornShuffleHandle<Integer, String, String> handle =
          new CelebornShuffleHandle<>(
              appId, host, port, userIdentifier, shuffleId, false, numMaps, dependency);
      final ShuffleClient client = new DummyShuffleClient(conf, tempFile);
      ((DummyShuffleClient) client).initReducePartitionMap(shuffleId, numPartitions, 1);
      File tempPusherFile = new File(tempDir, UUID.randomUUID().toString());
      SortBasedPusher pusher =
          createSortBasedPusher(conf, tempPusherFile, numPartitions, metrics.shuffleWriteMetrics());
      final SortBasedShuffleWriter<Integer, String, String> writer =
          createShuffleWriterWithPusher(
              handle, taskContext, conf, client, metrics.shuffleWriteMetrics(), pusher);
      AtomicInteger total = new AtomicInteger(0);
      System.out.println("time 1");
      Iterator iterator = getUnsafeRowIterator(32 * 1024, 8 * 1024, total, numPartitions);
      System.out.println(total.get());
      writer.doWrite(iterator);
      // given that the send buffer size is 32K and memory threshold is also 32K, with sortbasedwriter
      // we have pushed for 3 times (1 time per each existing partition),
      // so we will grow the threshold to 64K to mitigate too many pushes
      assertEquals(64 * 1024, pusher.getPushSortMemoryThreshold());
      System.out.println("time 2");
      iterator = getUnsafeRowIterator(32 * 1024, 8 * 1024, total, numPartitions);
      writer.doWrite(iterator);
      // no change on the threshold since we only pushed 32K more data
      assertEquals(64 * 1024, pusher.getPushSortMemoryThreshold());
      // write 32K more data will trigger the growth of threshold to 128K
      System.out.println("time 3");
      iterator = getUnsafeRowIterator(32 * 1024, 8 * 1024, total, numPartitions);
      writer.doWrite(iterator);
      assertEquals(128 * 1024, pusher.getPushSortMemoryThreshold());
      // some cleanup before the next test
      pusher.pushData(false);
      pusher.memoryThresholdManager.pushedCount = 0;
      pusher.memoryThresholdManager.pushedMemorySizeInBytes = 0;
      // now we insert 256K data but limited to 1 partition,
      // this time, we won't trigger the growing of memory threshold, because we have fully filled
      // the send buffer for each partition so that we are not "pushing data unnecessarily"
      System.out.println("time 4");
      iterator = getUnsafeRowIterator(256 * 1024, 8 * 1024, total, 1);
      writer.doWrite(iterator);
      assertEquals(128 * 1024, pusher.getPushSortMemoryThreshold());
    }
  }
  