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 static org.junit.Assert.assertNotNull;
  import static org.junit.Assert.assertTrue;
  import static org.junit.Assert.fail;
  
  import java.io.File;
  import java.io.FileInputStream;
  import java.io.IOException;
  import java.util.Collections;
  import java.util.Random;
  import java.util.UUID;
  import java.util.concurrent.atomic.AtomicInteger;
  
  import scala.None$;
  import scala.Option;
  import scala.Product2;
  import scala.Tuple2;
  import scala.collection.Iterator;
  import scala.collection.mutable.ListBuffer;
  
  import org.apache.spark.HashPartitioner;
  import org.apache.spark.Partitioner;
  import org.apache.spark.ShuffleDependency;
  import org.apache.spark.SparkConf;
  import org.apache.spark.SparkEnv;
  import org.apache.spark.TaskContext;
  import org.apache.spark.executor.ShuffleWriteMetrics;
  import org.apache.spark.executor.TaskMetrics;
  import org.apache.spark.memory.TaskMemoryManager;
  import org.apache.spark.memory.UnifiedMemoryManager;
  import org.apache.spark.scheduler.MapStatus;
  import org.apache.spark.serializer.KryoSerializer;
  import org.apache.spark.serializer.Serializer;
  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.PartitionIdPassthrough;
  import org.apache.spark.sql.execution.UnsafeRowSerializer;
  import org.apache.spark.sql.types.DataType;
  import org.apache.spark.sql.types.IntegerType$;
  import org.apache.spark.sql.types.StringType$;
  import org.apache.spark.storage.BlockManager;
  import org.apache.spark.storage.BlockManagerId;
  import org.apache.spark.unsafe.types.UTF8String;
  import org.junit.AfterClass;
  import org.junit.Before;
  import org.junit.BeforeClass;
  import org.junit.Test;
  import org.mockito.Answers;
  import org.mockito.Mock;
  import org.mockito.Mockito;
  import org.mockito.MockitoAnnotations;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import org.apache.celeborn.client.DummyShuffleClient;
  import org.apache.celeborn.client.ShuffleClient;
  import org.apache.celeborn.common.CelebornConf;
  import org.apache.celeborn.common.identity.UserIdentifier;
  import org.apache.celeborn.common.util.JavaUtils;
  import org.apache.celeborn.common.util.Utils;
  
  public abstract class CelebornShuffleWriterSuiteBase {
  
    private static final Logger LOG = LoggerFactory.getLogger(CelebornShuffleWriterSuiteBase.class);
    private static final String NORMAL_RECORD = "hello, world";
    private static final String GIANT_RECORD = getGiantRecord();
  
    private final Random rand = new Random(System.currentTimeMillis());
  
    private final String appId = "appId";
    private final String host = "host";
    private final int port = 0;
    private final int shuffleId = 0;
  
    private final UserIdentifier userIdentifier = new UserIdentifier("mock", "mock");
  
    private final int numMaps = 10;
    protected final int numPartitions = 10;
    private final SparkConf sparkConf = new SparkConf(false);
    private final BlockManagerId bmId = BlockManagerId.apply("execId", "host", 1, None$.empty());
  
    private final TaskMemoryManager taskMemoryManager =
        new TaskMemoryManager(UnifiedMemoryManager.apply(sparkConf, 1), 0);
    private final MapStatus mapStatus =
        SparkUtils.createMapStatus(bmId, new long[numPartitions], new long[numPartitions]);
  
    @Mock(answer = Answers.RETURNS_SMART_NULLS)
    private TaskContext taskContext;
  
    @Mock(answer = Answers.RETURNS_SMART_NULLS)
    private ShuffleDependency<Integer, String, String> dependency;
  
    @Mock(answer = Answers.RETURNS_SMART_NULLS)
    private SparkEnv env;
  
    @Mock(answer = Answers.RETURNS_SMART_NULLS)
    private BlockManager blockManager;
  
    private TaskMetrics metrics = null;
  
    private static File tempDir = null;
  
    public CelebornShuffleWriterSuiteBase() throws IOException {}
  
    @BeforeClass
    public static void beforeAll() {
      tempDir = Utils.createTempDir(System.getProperty("java.io.tmpdir"), "celeborn_test");
    }
  
    @AfterClass
    public static void afterAll() {
      try {
        JavaUtils.deleteRecursively(tempDir);
      } catch (IOException e) {
        LOG.error("Failed to delete temporary directory.", e);
      }
    }
  
    @Before
    public void beforeEach() {
      metrics = new TaskMetrics();
      MockitoAnnotations.initMocks(this);
      Mockito.doReturn(shuffleId).when(dependency).shuffleId();
  
      Mockito.doReturn(metrics).when(taskContext).taskMetrics();
      Mockito.doReturn(taskMemoryManager).when(taskContext).taskMemoryManager();
  
      Mockito.doReturn(bmId).when(blockManager).shuffleServerId();
      Mockito.doReturn(blockManager).when(env).blockManager();
      Mockito.doReturn(sparkConf).when(env).conf();
      SparkEnv.set(env);
    }
  
    @Test
    public void testEmptyBlock() throws Exception {
      final KryoSerializer serializer = new KryoSerializer(sparkConf);
      final CelebornConf conf = new CelebornConf();
      check(0, conf, serializer);
    }
  
    @Test
    public void testEmptyBlockWithFastWrite() throws Exception {
      final UnsafeRowSerializer serializer = new UnsafeRowSerializer(2, null);
      final CelebornConf conf = new CelebornConf();
      check(0, conf, serializer);
    }
  
    @Test
    public void testMergeSmallBlock() throws Exception {
      final KryoSerializer serializer = new KryoSerializer(sparkConf);
      final CelebornConf conf =
          new CelebornConf().set(CelebornConf.CLIENT_PUSH_BUFFER_MAX_SIZE().key(), "1024");
      check(10000, conf, serializer);
    }
  
    @Test
    public void testMergeSmallBlockWithFastWrite() throws Exception {
      final UnsafeRowSerializer serializer = new UnsafeRowSerializer(2, null);
      final CelebornConf conf =
          new CelebornConf().set(CelebornConf.CLIENT_PUSH_BUFFER_MAX_SIZE().key(), "1024");
      check(10000, conf, serializer);
    }
  
    @Test
    public void testGiantRecord() throws Exception {
      final KryoSerializer serializer = new KryoSerializer(sparkConf);
      final CelebornConf conf =
          new CelebornConf().set(CelebornConf.CLIENT_PUSH_BUFFER_MAX_SIZE().key(), "5");
      check(10000, conf, serializer);
    }
  
    @Test
    public void testGiantRecordWithFastWrite() throws Exception {
      final UnsafeRowSerializer serializer = new UnsafeRowSerializer(2, null);
      final CelebornConf conf =
          new CelebornConf().set(CelebornConf.CLIENT_PUSH_BUFFER_MAX_SIZE().key(), "5");
      check(10000, conf, serializer);
    }
  
    @Test
    public void testGiantRecordAndMergeSmallBlock() throws Exception {
      final KryoSerializer serializer = new KryoSerializer(sparkConf);
      final CelebornConf conf =
          new CelebornConf().set(CelebornConf.CLIENT_PUSH_BUFFER_MAX_SIZE().key(), "128");
      check(2 << 30, conf, serializer);
    }
  
    @Test
    public void testGiantRecordAndMergeSmallBlockWithFastWrite() throws Exception {
      final UnsafeRowSerializer serializer = new UnsafeRowSerializer(2, null);
      final CelebornConf conf =
          new CelebornConf().set(CelebornConf.CLIENT_PUSH_BUFFER_MAX_SIZE().key(), "128");
      check(2 << 30, conf, serializer);
    }
  
    private void check(
        final int approximateSize, final CelebornConf conf, final Serializer serializer)
        throws Exception {
      final boolean useUnsafe = serializer instanceof UnsafeRowSerializer;
  
      final Partitioner partitioner =
          useUnsafe ? new PartitionIdPassthrough(numPartitions) : new HashPartitioner(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);
  
      final ShuffleWriter<Integer, String> writer =
          createShuffleWriter(handle, taskContext, conf, client);
  
      if (writer instanceof SortBasedShuffleWriter) {
        assertEquals(useUnsafe, ((SortBasedShuffleWriter) writer).canUseFastWrite());
      } else if (writer instanceof HashBasedShuffleWriter) {
        assertEquals(useUnsafe, ((HashBasedShuffleWriter) writer).canUseFastWrite());
      }
  
      AtomicInteger total = new AtomicInteger(0);
      Iterator iterator = getIterator(approximateSize, total, useUnsafe, false);
  
      int expectChecksum = 0;
      for (int i = 0; i < total.intValue(); ++i) {
        expectChecksum ^= i;
      }
  
      writer.write(iterator);
      Option<MapStatus> status = writer.stop(true);
      client.shutdown();
  
      assertNotNull(status);
      assertTrue(status.isDefined());
      assertEquals(bmId, status.get().location());
  
      ShuffleWriteMetrics metrics = taskContext.taskMetrics().shuffleWriteMetrics();
      assertEquals(metrics.recordsWritten(), total.intValue());
      assertEquals(metrics.bytesWritten(), tempFile.length());
  
      try (FileInputStream fis = new FileInputStream(tempFile)) {
        Iterator it = serializer.newInstance().deserializeStream(fis).asKeyValueIterator();
        int checksum = 0;
        while (it.hasNext()) {
          Product2<Integer, ?> record;
          if (useUnsafe) {
            record = (Product2<Integer, UnsafeRow>) it.next();
          } else {
            record = (Product2<Integer, String>) it.next();
          }
  
          assertNotNull(record);
          assertNotNull(record._1());
          assertNotNull(record._2());
  
          int key;
          String value;
  
          if (useUnsafe) {
            UnsafeRow row = (UnsafeRow) record._2();
  
            key = row.getInt(0);
            value = row.getString(1);
          } else {
            key = record._1();
            value = (String) record._2();
          }
  
          checksum ^= key;
          total.decrementAndGet();
  
          assertTrue(
              "value should equals to normal record or giant record with key.",
              value.equals(key + ": " + NORMAL_RECORD) || value.equals(key + ": " + GIANT_RECORD));
        }
        assertEquals(0, total.intValue());
        assertEquals(expectChecksum, checksum);
      } catch (Exception e) {
        e.printStackTrace();
        fail("Should read with no exception.");
      }
    }
  
    private Iterator getIterator(
        final int size, final AtomicInteger total, final boolean unsafe, final boolean mix) {
      if (unsafe) {
        return getUnsafeRowIterator(size, total, mix);
      } else {
        return getNormalIterator(size, total, mix);
      }
    }
  
    private Iterator<Product2<Integer, String>> getNormalIterator(
        final int size, final AtomicInteger total, final boolean mix) {
      int current = 0;
      ListBuffer<Product2<Integer, String>> list = new ListBuffer<>();
      while (current < size) {
        int key = total.getAndIncrement();
        String value = key + ": " + (mix && rand.nextBoolean() ? GIANT_RECORD : NORMAL_RECORD);
        current += value.length();
        list.$plus$eq(new Tuple2<>(key, value));
      }
      return list.toIterator();
    }
  
    private Iterator<Product2<Integer, UnsafeRow>> getUnsafeRowIterator(
        final int size, final AtomicInteger total, final boolean mix) {
      int current = 0;
      ListBuffer<Product2<Integer, UnsafeRow>> list = new ListBuffer<>();
      while (current < size) {
        int key = total.getAndIncrement();
        String value = key + ": " + (mix && rand.nextBoolean() ? GIANT_RECORD : NORMAL_RECORD);
        current += value.length();
  
        ListBuffer<Object> values = new ListBuffer<>();
        values.$plus$eq(key);
        values.$plus$eq(UTF8String.fromString(value));
  
        InternalRow row = InternalRow.apply(values.toSeq());
        DataType[] types = new DataType[2];
        types[0] = IntegerType$.MODULE$;
        types[1] = StringType$.MODULE$;
        UnsafeRow unsafeRow = UnsafeProjection.create(types).apply(row);
  
        list.$plus$eq(new Tuple2<>(key % numPartitions, unsafeRow));
      }
      return list.toIterator();
    }
  
    private static String getGiantRecord() {
      int numCopies = (128 + NORMAL_RECORD.length() - 1) / NORMAL_RECORD.length();
      return String.join("/", Collections.nCopies(numCopies, NORMAL_RECORD));
    }
  
    protected abstract ShuffleWriter<Integer, String> createShuffleWriter(
        CelebornShuffleHandle handle, TaskContext context, CelebornConf conf, ShuffleClient client)
        throws IOException;
  }
  