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.celeborn.service.deploy.worker.storage;
  
  import java.io.IOException;
  import java.util.Objects;
  import java.util.concurrent.atomic.AtomicInteger;
  
  import scala.Option;
  
  import com.google.common.annotations.VisibleForTesting;
  import com.google.protobuf.GeneratedMessageV3;
  import io.netty.buffer.ByteBuf;
  import org.roaringbitmap.RoaringBitmap;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import org.apache.celeborn.common.CelebornConf;
  import org.apache.celeborn.common.meta.*;
  import org.apache.celeborn.common.metrics.source.AbstractSource;
  import org.apache.celeborn.common.protocol.PartitionSplitMode;
  import org.apache.celeborn.common.protocol.StorageInfo;
  import org.apache.celeborn.service.deploy.worker.WorkerSource;
  import org.apache.celeborn.service.deploy.worker.congestcontrol.CongestionController;
  import org.apache.celeborn.service.deploy.worker.congestcontrol.UserCongestionControlContext;
  import org.apache.celeborn.service.deploy.worker.memory.MemoryManager;
  
  /*
   * Note: Once FlushNotifier.exception is set, the whole file is not available.
   *       That's fine some of the internal state(e.g. bytesFlushed) may be inaccurate.
   */
  public class PartitionDataWriter implements DeviceObserver {
    private static final Logger logger = LoggerFactory.getLogger(PartitionDataWriter.class);
  
    private final long splitThreshold;
    private final PartitionSplitMode splitMode;
    private final long memoryFileStorageMaxFileSize;
    private final AtomicInteger numPendingWrites = new AtomicInteger(0);
    private final PartitionDataWriterContext writerContext;
    protected final AbstractSource source; // metrics
    private final String writerString;
    private final StorageManager storageManager;
    private final FlushNotifier notifier = new FlushNotifier();
  
    private UserCongestionControlContext userCongestionControlContext = null;
    private volatile TierWriterBase currentTierWriter;
  
    public PartitionDataWriter(
        StorageManager storageManager,
        AbstractSource workerSource,
        CelebornConf conf,
        DeviceMonitor deviceMonitor,
        PartitionDataWriterContext writerContext) {
      memoryFileStorageMaxFileSize = conf.workerMemoryFileStorageMaxFileSize();
      this.writerContext = writerContext;
      this.source = workerSource;
  
      this.storageManager = storageManager;
      this.splitThreshold = writerContext.getSplitThreshold();
      this.splitMode = writerContext.getPartitionSplitMode();
      String shuffleKey = writerContext.getShuffleKey();
      String filename = writerContext.getPartitionLocation().getFileName();
  
      writerString = shuffleKey + "-" + filename + "-partition-writer";
  
      logger.debug("FileWriter {} split threshold {} mode {}", this, splitThreshold, splitMode);
      if (CongestionController.instance() != null) {
        userCongestionControlContext =
            CongestionController.instance()
                .getUserCongestionContext(writerContext.getUserIdentifier());
      }
  
      writerContext.setPartitionDataWriter(this);
      writerContext.setDeviceMonitor(deviceMonitor);
      currentTierWriter =
          storageManager.storagePolicy().createFileWriter(writerContext, numPendingWrites, notifier);
    }
  
    public DiskFileInfo getDiskFileInfo() {
      // keep compatible with current logic
      FileInfo currentFileInfo = currentTierWriter.fileInfo();
      if (currentFileInfo instanceof DiskFileInfo) {
        return (DiskFileInfo) currentFileInfo;
      }
      return null;
    }
  
    public String getFilePath() {
      DiskFileInfo diskFileInfo = getDiskFileInfo();
      if (diskFileInfo != null) {
        return diskFileInfo.getFilePath();
      }
      return "";
    }
  
    public void incrementPendingWrites() {
      numPendingWrites.incrementAndGet();
    }
  
    public void decrementPendingWrites() {
      numPendingWrites.decrementAndGet();
    }
  
    // this will only happen if a worker is under memory pressure
    @VisibleForTesting
    public synchronized void flush() {
      currentTierWriter.flush(false, false);
    }
  
    public synchronized boolean needHardSplitForMemoryShuffleStorage() {
      if (!(currentTierWriter instanceof MemoryTierWriter)) {
        return false;
      }
      return !storageManager.localOrDfsStorageAvailable()
          && (currentTierWriter.fileInfo().getFileLength() > memoryFileStorageMaxFileSize
              || !MemoryManager.instance().memoryFileStorageAvailable());
    }
  
    public synchronized void write(ByteBuf data) throws IOException {
      if (currentTierWriter.needEvict()) {
        evict(false);
      }
      currentTierWriter.write(data);
    }
  
    public RoaringBitmap getMapIdBitMap() {
      Option<RoaringBitmap> bitmapOpt = currentTierWriter.metaHandler().getMapIdBitmap();
      // to keep compatible with scala 2.11
      if (bitmapOpt.isDefined()) {
        return bitmapOpt.get();
      }
      return null;
    }
  
    public StorageInfo getStorageInfo() {
      StorageInfo storageInfo = null;
      FileInfo fileInfo = currentTierWriter.fileInfo();
      if (fileInfo instanceof DiskFileInfo) {
        DiskFileInfo diskFileInfo = (DiskFileInfo) fileInfo;
        if (diskFileInfo.isDFS()) {
          if (((DfsTierWriter) currentTierWriter).deleted()) {
            return null;
          } else if (diskFileInfo.isS3()) {
            storageInfo = new StorageInfo(StorageInfo.Type.S3, true, diskFileInfo.getFilePath());
          } else if (diskFileInfo.isOSS()) {
            storageInfo = new StorageInfo(StorageInfo.Type.OSS, true, diskFileInfo.getFilePath());
          } else {
            storageInfo = new StorageInfo(StorageInfo.Type.HDFS, true, diskFileInfo.getFilePath());
          }
        } else {
          LocalFlusher flusher = (LocalFlusher) currentTierWriter.getFlusher();
          storageInfo = new StorageInfo(flusher.diskType(), true, "");
        }
      } else if (fileInfo instanceof MemoryFileInfo) {
        storageInfo = new StorageInfo(StorageInfo.Type.MEMORY, true, "");
      }
      if (storageInfo != null
          && currentTierWriter.fileInfo().getFileMeta() instanceof ReduceFileMeta) {
        storageInfo.setFileSize(currentTierWriter.fileInfo().getFileLength());
        storageInfo.setChunkOffsets(
            ((ReduceFileMeta) currentTierWriter.fileInfo().getFileMeta()).getChunkOffsets());
      }
      return storageInfo;
    }
  
    public boolean isClosed() {
      return currentTierWriter.closed();
    }
  
    // evict and flush method need to be in a same synchronized block
    // because memory manager may want to evict a file under memory pressure
    public synchronized void evict(boolean checkClose) {
      // close and evict might be invoked concurrently
      // do not evict committed files from memory manager
      // evict memory file info if worker is shutdown gracefully
      if (checkClose) {
        if (currentTierWriter.closed()) {
          return;
        }
      }
      TierWriterBase newTierWriter =
          storageManager
              .storagePolicy()
              .getEvictedFileWriter(currentTierWriter, writerContext, numPendingWrites, notifier);
      currentTierWriter.evict(newTierWriter);
      currentTierWriter = newTierWriter;
    }
  
    public synchronized void destroy(IOException ioException) {
      currentTierWriter.destroy(ioException);
    }
  
    public synchronized long close() {
      long length = currentTierWriter.close();
      source.updateHistogram(WorkerSource.PARTITION_FILE_SIZE(), length);
      return length;
    }
  
    public FileInfo getCurrentFileInfo() {
      return currentTierWriter.fileInfo();
    }
  
    public IOException getException() {
      if (notifier.hasException()) {
        return notifier.exception.get();
      } else {
        return null;
      }
    }
  
    @Override
    public boolean equals(Object o) {
      if (o == null || getClass() != o.getClass()) return false;
      PartitionDataWriter that = (PartitionDataWriter) o;
      return Objects.equals(writerString, that.writerString);
    }
  
    @Override
    public int hashCode() {
      return Objects.hashCode(writerString);
    }
  
    @Override
    public String toString() {
      return writerString;
    }
  
    public void flushOnMemoryPressure() throws IOException {
      // this won't happen if this writer is in memory
      flush();
    }
  
    public long getSplitThreshold() {
      return splitThreshold;
    }
  
    public PartitionSplitMode getSplitMode() {
      return splitMode;
    }
  
    @Override
    public void notifyError(String mountPoint, DiskStatus diskStatus) {
      destroy(
          new IOException(
              "Destroy FileWriter "
                  + this
                  + " by device ERROR."
                  + " Disk: "
                  + mountPoint
                  + " Status: "
                  + diskStatus));
    }
  
    // These empty methods are intended to match scala 2.11 restrictions that
    // trait can not be used as an interface with default implementation.
    @Override
    public void notifyHealthy(String mountPoint) {}
  
    @Override
    public void notifyHighDiskUsage(String mountPoint) {}
  
    @Override
    public void notifyNonCriticalError(String mountPoint, DiskStatus diskStatus) {}
  
    public MemoryFileInfo getMemoryFileInfo() {
      return ((MemoryFileInfo) currentTierWriter.fileInfo());
    }
  
    public UserCongestionControlContext getUserCongestionControlContext() {
      return userCongestionControlContext;
    }
  
    public void handleEvents(GeneratedMessageV3 message) {
      currentTierWriter.metaHandler().handleEvent(message);
    }
  
    public PartitionMetaHandler getMetaHandler() {
      return currentTierWriter.metaHandler();
    }
  
    public Flusher getFlusher() {
      return currentTierWriter.getFlusher();
    }
  }
  