<commit_before>// Copyright 2022 Google LLC
//
// Licensed 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 com.google.cloud.spanner.pgadapter.utils;

import com.google.api.core.ApiFuture;
import com.google.api.core.ApiFutures;
import com.google.api.core.InternalApi;
import com.google.api.core.SettableApiFuture;
import com.google.cloud.ByteArray;
import com.google.cloud.spanner.DatabaseClient;
import com.google.cloud.spanner.ErrorCode;
import com.google.cloud.spanner.Mutation;
import com.google.cloud.spanner.Mutation.WriteBuilder;
import com.google.cloud.spanner.SpannerException;
import com.google.cloud.spanner.SpannerExceptionFactory;
import com.google.cloud.spanner.Type;
import com.google.cloud.spanner.Value;
import com.google.cloud.spanner.connection.Connection;
import com.google.cloud.spanner.connection.StatementResult;
import com.google.cloud.spanner.pgadapter.parsers.copy.CopyTreeParser.CopyOptions;
import com.google.cloud.spanner.pgadapter.statements.BackendConnection.UpdateCount;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.collect.ImmutableList;
import com.google.common.util.concurrent.FutureCallback;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.common.util.concurrent.ListeningExecutorService;
import com.google.common.util.concurrent.MoreExecutors;
import java.io.Closeable;
import java.io.IOException;
import java.io.PipedOutputStream;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.concurrent.Callable;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Executors;
import java.util.concurrent.LinkedBlockingDeque;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.logging.Level;
import java.util.logging.Logger;
import javax.annotation.Nonnull;
import javax.annotation.concurrent.GuardedBy;
import org.apache.commons.csv.CSVFormat;

@InternalApi
public class MutationWriter implements Callable<StatementResult>, Closeable {
  public enum CopyTransactionMode {
    /**
     * 'Normal' auto-commit mode. The entire COPY operation is atomic. If the number of mutations
     * exceeds any of the transaction limits of Cloud Spanner, the COPY operation will fail.
     */
    ImplicitAtomic,
    /**
     * The COPY operation is executed as a series of (partly parallel) transactions. The COPY
     * operation is not atomic, and an error halfway the operation can mean that some data was
     * persisted to the database.
     */
    ImplicitNonAtomic,
    /**
     * There is an explicit transaction on the connection. The COPY will be part of that
     * transaction.
     */
    Explicit,
  }

  private static final Logger logger = Logger.getLogger(MutationWriter.class.getName());

  private static final int DEFAULT_MUTATION_LIMIT = 20_000; // 20k mutation count limit
  private static final int DEFAULT_COMMIT_LIMIT =
      100_000_000; // 100MB mutation API commit size limit
  /**
   * The factor that the size of the incoming payload is multiplied with to estimate whether the
   * current commit size will exceed the maximum commit size. The factor is chosen conservatively to
   * ensure that a COPY operation of a large file does not unnecessarily break because of an
   * exceeded commit size.
   */
  private static final float DEFAULT_COMMIT_LIMIT_MULTIPLIER_FACTOR = 2.0f;

  private static final int DEFAULT_MAX_PARALLELISM = 128;
  private static final int DEFAULT_PIPE_BUFFER_SIZE = 1 << 16;

  /**
   * COPY will INSERT records by default. This is consistent with how COPY on PostgreSQL works. This
   * option allows PGAdapter to use InsertOrUpdate instead. This can be slightly more efficient for
   * bulk uploading, and it makes it easier to retry a failed non-atomic batch that might have
   * already uploaded some but not all data.
   */
  private final boolean insertOrUpdate =
      Boolean.parseBoolean(System.getProperty("copy_in_insert_or_update", "false"));

  private final int commitSizeLimit =
      Integer.parseInt(
          System.getProperty("copy_in_commit_limit", String.valueOf(DEFAULT_COMMIT_LIMIT)));
  private final CopyTransactionMode transactionMode;
  private long rowCount;
  private final Connection connection;
  private final String tableName;
  private final Map<String, Type> tableColumns;
  private final int maxBatchSize;
  private final long commitSizeLimitForBatching;
  private final int maxParallelism;
  private final CopyInParser parser;
  private final PipedOutputStream payload = new PipedOutputStream();
  private final AtomicBoolean commit = new AtomicBoolean(false);
  private final AtomicBoolean rollback = new AtomicBoolean(false);
  private final CountDownLatch closedLatch = new CountDownLatch(1);
  private final ListeningExecutorService executorService =
      MoreExecutors.listeningDecorator(Executors.newCachedThreadPool());

  private final Object lock = new Object();

  @GuardedBy("lock")
  private SpannerException exception;

  public MutationWriter(
      CopyTransactionMode transactionMode,
      Connection connection,
      String tableName,
      Map<String, Type> tableColumns,
      int indexedColumnsCount,
      CopyOptions.Format copyFormat,
      CSVFormat format,
      boolean hasHeader)
      throws IOException {
    this.transactionMode = transactionMode;
    this.connection = connection;
    this.tableName = tableName;
    this.tableColumns = tableColumns;
    int mutationLimit =
        Math.max(
            Integer.parseInt(
                System.getProperty(
                    "copy_in_mutation_limit", String.valueOf(DEFAULT_MUTATION_LIMIT))),
            1);
    this.maxBatchSize = Math.max(mutationLimit / (tableColumns.size() + indexedColumnsCount), 1);
    float commitLimitMultiplierFactor =
        Math.max(
            Float.parseFloat(
                System.getProperty(
                    "copy_in_commit_limit_multiplier_factor",
                    String.valueOf(DEFAULT_COMMIT_LIMIT_MULTIPLIER_FACTOR))),
            1.0f);
    this.commitSizeLimitForBatching =
        Math.round((float) commitSizeLimit / commitLimitMultiplierFactor);
    this.maxParallelism =
        Math.max(
            Integer.parseInt(
                System.getProperty(
                    "copy_in_max_parallelism", String.valueOf(DEFAULT_MAX_PARALLELISM))),
            1);
    int pipeBufferSize =
        Math.max(
            Integer.parseInt(
                System.getProperty(
                    "copy_in_pipe_buffer_size", String.valueOf(DEFAULT_PIPE_BUFFER_SIZE))),
            1024);
    this.parser = CopyInParser.create(copyFormat, format, payload, pipeBufferSize, hasHeader);
  }

  /** @return number of rows copied into Spanner */
  public long getRowCount() {
    return this.rowCount;
  }

  public void addCopyData(byte[] payload) {
    synchronized (lock) {
      if (this.exception != null) {
        throw this.exception;
      }
    }
    try {
      this.payload.write(payload);
    } catch (IOException e) {
      // Ignore the exception if the executor has already been shutdown. That means that an error
      // occurred that ended the COPY operation while we were writing data to the buffer.
      if (!executorService.isShutdown()) {
        SpannerException spannerException =
            SpannerExceptionFactory.newSpannerException(
                ErrorCode.INTERNAL, "Could not write copy data to buffer", e);
        logger.log(Level.SEVERE, spannerException.getMessage(), spannerException);
        throw spannerException;
      }
    }
  }

  /** Indicate that this mutation writer should commit. */
  public void commit() {
    this.commit.set(true);
  }

  /**
   * Indicate that this mutation writer should be rolled back. This will not rollback any changes
   * that have already been committed if the mutation writer is running in {@link
   * CopyTransactionMode#ImplicitNonAtomic}.
   */
  public void rollback() {
    this.rollback.set(true);
  }

  @Override
  public void close() throws IOException {
    this.payload.close();
    this.closedLatch.countDown();
  }

  @Override
  public StatementResult call() throws Exception {
    // This LinkedBlockingDeque holds a reference to all transactions that are currently active. The
    // max capacity of this deque is what ensures that we never have more than maxParallelism
    // transactions running at the same time. We could also achieve that by using a thread pool with
    // a fixed number of threads. The problem with that is however that Java does not have a thread
    // pool implementation that will block if a new task is offered and all threads are currently in
    // use. The only options are 'fail or add to queue'. We want to block our worker thread in this
    // case when the max parallel transactions has been reached, as that automatically creates back-
    // pressure in our entire pipeline that consists of:
    // Client app (psql) -> CopyData message -> CSVParser -> Transaction.
    LinkedBlockingDeque<ApiFuture<Void>> activeCommitFutures =
        new LinkedBlockingDeque<>(maxParallelism);
    // This list holds all transactions that we have started. We will wait on this entire list
    // before finishing, to ensure that all data has been written before we signal that we are done.
    List<ApiFuture<Void>> allCommitFutures = new ArrayList<>();
    try {
      Iterator<CopyRecord> iterator = this.parser.iterator();
      List<Mutation> mutations = new ArrayList<>();
      long currentBufferByteSize = 0L;
      // Note: iterator.hasNext() blocks if there is not enough data in the pipeline to construct a
      // complete record. It returns false if the stream has been closed and all records have been
      // returned.
      while (!rollback.get() && iterator.hasNext()) {
        CopyRecord record = iterator.next();
        if (record.numColumns() != this.tableColumns.keySet().size()) {
          throw SpannerExceptionFactory.newSpannerException(
              ErrorCode.INVALID_ARGUMENT,
              "Invalid COPY data: Row length mismatched. Expected "
                  + this.tableColumns.keySet().size()
                  + " columns, but only found "
                  + record.numColumns());
        }

        Mutation mutation = buildMutation(record);
        int mutationSize = calculateSize(mutation);
        this.rowCount++;

        if (transactionMode == CopyTransactionMode.ImplicitNonAtomic) {
          currentBufferByteSize =
              addMutationAndMaybeFlushTransaction(
                  activeCommitFutures,
                  allCommitFutures,
                  mutations,
                  mutation,
                  currentBufferByteSize,
                  mutationSize);
        } else {
          mutations.add(mutation);
          currentBufferByteSize += mutationSize;
          if (mutations.size() > maxBatchSize) {
            throw SpannerExceptionFactory.newSpannerException(
                ErrorCode.FAILED_PRECONDITION,
                "Record count: "
                    + mutations.size()
                    + " has exceeded the limit: "
                    + maxBatchSize
                    + ".\n\nThe number of mutations per record is equal to the number of columns in the record "
                    + "plus the number of indexed columns in the record. The maximum number of mutations "
                    + "in one transaction is "
                    + DEFAULT_MUTATION_LIMIT
                    + ".\n\nExecute `SET SPANNER.AUTOCOMMIT_DML_MODE='PARTITIONED_NON_ATOMIC'` before executing a large COPY operation "
                    + "to instruct PGAdapter to automatically break large transactions into multiple smaller. "
                    + "This will make the COPY operation non-atomic.\n\n");
          }
          if (currentBufferByteSize > commitSizeLimit) {
            throw SpannerExceptionFactory.newSpannerException(
                ErrorCode.FAILED_PRECONDITION,
                "Commit size: "
                    + currentBufferByteSize
                    + " has exceeded the limit: "
                    + commitSizeLimit
                    + ".\n\nExecute `SET SPANNER.AUTOCOMMIT_DML_MODE='PARTITIONED_NON_ATOMIC'` before executing a large COPY operation "
                    + "to instruct PGAdapter to automatically break large transactions into multiple smaller. "
                    + "This will make the COPY operation non-atomic.\n\n");
          }
        }
      } // end of iterator.hasNext()

      // There are no more CSVRecords in the pipeline.
      // Write any remaining mutations in the buffer.
      if (!rollback.get() && !mutations.isEmpty()) {
        if (transactionMode == CopyTransactionMode.Explicit) {
          connection.write(mutations);
        } else {
          // Wait until we have received a CopyDone message before writing the remaining data to
          // Spanner. If we are in a non-atomic transaction, there might already be data that have
          // been written to Spanner.
          closedLatch.await();
          if (commit.get()) {
            allCommitFutures.add(writeToSpannerAsync(activeCommitFutures, mutations));
          }
        }
      }
      // Wait for all commits to finish. We do this even if something went wrong, as it ensures two
      // things:
      // 1. All commits that were in flight when something went wrong will finish before we return
      //    an error to the client application. This prevents commits still being added to the
      //    database after we have returned an error, which could cause confusion.
      // 2. This will throw the underlying exception, so we can catch and register it.
      ApiFutures.allAsList(allCommitFutures).get();
    } catch (SpannerException e) {
      synchronized (lock) {
        this.exception = e;
        throw this.exception;
      }
    } catch (ExecutionException e) {
      synchronized (lock) {
        this.exception = SpannerExceptionFactory.asSpannerException(e.getCause());
        throw this.exception;
      }
    } catch (Exception e) {
      synchronized (lock) {
        this.exception = SpannerExceptionFactory.asSpannerException(e);
        throw this.exception;
      }
    } finally {
      this.executorService.shutdown();
      if (!this.executorService.awaitTermination(60L, TimeUnit.SECONDS)) {
        logger.log(Level.WARNING, "Timeout while waiting for MutationWriter executor to shutdown.");
      }
      this.payload.close();
      this.parser.close();
    }
    return new UpdateCount(rowCount);
  }

  private long addMutationAndMaybeFlushTransaction(
      LinkedBlockingDeque<ApiFuture<Void>> activeCommitFutures,
      List<ApiFuture<Void>> allCommitFutures,
      List<Mutation> mutations,
      Mutation mutation,
      long currentBufferByteSize,
      int mutationSize)
      throws Exception {
    // Flush before adding to the buffer if the commit size would be exceeded.
    if (!mutations.isEmpty() && currentBufferByteSize + mutationSize > commitSizeLimitForBatching) {
      allCommitFutures.add(writeToSpannerAsync(activeCommitFutures, mutations));
      mutations.clear();
      mutations.add(mutation);
      return mutationSize;
    }

    mutations.add(mutation);
    if (mutations.size() == maxBatchSize) {
      allCommitFutures.add(writeToSpannerAsync(activeCommitFutures, mutations));
      mutations.clear();
      return 0L; // Buffer is empty, so the batch size in bytes is now back to zero.
    }
    return currentBufferByteSize + mutationSize;
  }

  private ApiFuture<Void> writeToSpannerAsync(
      LinkedBlockingDeque<ApiFuture<Void>> activeCommitFutures, Iterable<Mutation> mutations)
      throws Exception {

    SettableApiFuture<Void> settableApiFuture = SettableApiFuture.create();
    // Add this future to the list of active commit futures. This will block if the deque is full,
    // and this will effectively apply back-pressure to the entire stream as the worker thread is
    // blocked until there is room in the deque.
    activeCommitFutures.put(settableApiFuture);

    DatabaseClient dbClient = connection.getDatabaseClient();
    ImmutableList<Mutation> immutableMutations = ImmutableList.copyOf(mutations);
    ListenableFuture<Void> listenableFuture =
        executorService.submit(
            () -> {
              dbClient.write(immutableMutations);
              return null;
            });
    Futures.addCallback(
        listenableFuture,
        new FutureCallback<Void>() {
          @Override
          public void onFailure(@Nonnull Throwable t) {
            rollback.set(true);
            //noinspection ResultOfMethodCallIgnored
            activeCommitFutures.remove(settableApiFuture);
            settableApiFuture.setException(t);
          }

          @Override
          public void onSuccess(Void result) {
            //noinspection ResultOfMethodCallIgnored
            activeCommitFutures.remove(settableApiFuture);<commit_msg>fix: use a longer timeout for Commit for COPY (#308)<commit_after>// Copyright 2022 Google LLC
//
// Licensed 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 com.google.cloud.spanner.pgadapter.utils;

import com.google.api.core.ApiFuture;
import com.google.api.core.ApiFutures;
import com.google.api.core.InternalApi;
import com.google.api.core.SettableApiFuture;
import com.google.cloud.ByteArray;
import com.google.cloud.spanner.DatabaseClient;
import com.google.cloud.spanner.ErrorCode;
import com.google.cloud.spanner.Mutation;
import com.google.cloud.spanner.Mutation.WriteBuilder;
import com.google.cloud.spanner.SpannerException;
import com.google.cloud.spanner.SpannerExceptionFactory;
import com.google.cloud.spanner.SpannerOptions;
import com.google.cloud.spanner.SpannerOptions.SpannerCallContextTimeoutConfigurator;
import com.google.cloud.spanner.Type;
import com.google.cloud.spanner.Value;
import com.google.cloud.spanner.connection.Connection;
import com.google.cloud.spanner.connection.StatementResult;
import com.google.cloud.spanner.pgadapter.parsers.copy.CopyTreeParser.CopyOptions;
import com.google.cloud.spanner.pgadapter.statements.BackendConnection.UpdateCount;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.collect.ImmutableList;
import com.google.common.util.concurrent.FutureCallback;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.common.util.concurrent.ListeningExecutorService;
import com.google.common.util.concurrent.MoreExecutors;
import io.grpc.Context;
import java.io.Closeable;
import java.io.IOException;
import java.io.PipedOutputStream;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.concurrent.Callable;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Executors;
import java.util.concurrent.LinkedBlockingDeque;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.logging.Level;
import java.util.logging.Logger;
import javax.annotation.Nonnull;
import javax.annotation.concurrent.GuardedBy;
import org.apache.commons.csv.CSVFormat;
import org.threeten.bp.Duration;

@InternalApi
public class MutationWriter implements Callable<StatementResult>, Closeable {
  public enum CopyTransactionMode {
    /**
     * 'Normal' auto-commit mode. The entire COPY operation is atomic. If the number of mutations
     * exceeds any of the transaction limits of Cloud Spanner, the COPY operation will fail.
     */
    ImplicitAtomic,
    /**
     * The COPY operation is executed as a series of (partly parallel) transactions. The COPY
     * operation is not atomic, and an error halfway the operation can mean that some data was
     * persisted to the database.
     */
    ImplicitNonAtomic,
    /**
     * There is an explicit transaction on the connection. The COPY will be part of that
     * transaction.
     */
    Explicit,
  }

  private static final Logger logger = Logger.getLogger(MutationWriter.class.getName());

  private static final int DEFAULT_MUTATION_LIMIT = 20_000; // 20k mutation count limit
  private static final int DEFAULT_COMMIT_LIMIT =
      100_000_000; // 100MB mutation API commit size limit
  /**
   * The factor that the size of the incoming payload is multiplied with to estimate whether the
   * current commit size will exceed the maximum commit size. The factor is chosen conservatively to
   * ensure that a COPY operation of a large file does not unnecessarily break because of an
   * exceeded commit size.
   */
  private static final float DEFAULT_COMMIT_LIMIT_MULTIPLIER_FACTOR = 2.0f;

  private static final int DEFAULT_COMMIT_TIMEOUT_SECONDS = 300;
  private static final int DEFAULT_MAX_PARALLELISM = 128;
  private static final int DEFAULT_PIPE_BUFFER_SIZE = 1 << 16;

  /**
   * COPY will INSERT records by default. This is consistent with how COPY on PostgreSQL works. This
   * option allows PGAdapter to use InsertOrUpdate instead. This can be slightly more efficient for
   * bulk uploading, and it makes it easier to retry a failed non-atomic batch that might have
   * already uploaded some but not all data.
   */
  private final boolean insertOrUpdate =
      Boolean.parseBoolean(System.getProperty("copy_in_insert_or_update", "false"));

  private final int commitSizeLimit =
      Integer.parseInt(
          System.getProperty("copy_in_commit_limit", String.valueOf(DEFAULT_COMMIT_LIMIT)));
  private final int commitTimeoutSeconds =
      Integer.parseInt(
          System.getProperty(
              "copy_in_commit_timeout_seconds", String.valueOf(DEFAULT_COMMIT_TIMEOUT_SECONDS)));
  private final CopyTransactionMode transactionMode;
  private long rowCount;
  private final Connection connection;
  private final String tableName;
  private final Map<String, Type> tableColumns;
  private final int maxBatchSize;
  private final long commitSizeLimitForBatching;
  private final int maxParallelism;
  private final CopyInParser parser;
  private final PipedOutputStream payload = new PipedOutputStream();
  private final AtomicBoolean commit = new AtomicBoolean(false);
  private final AtomicBoolean rollback = new AtomicBoolean(false);
  private final CountDownLatch closedLatch = new CountDownLatch(1);
  private final ListeningExecutorService executorService =
      MoreExecutors.listeningDecorator(Executors.newCachedThreadPool());

  private final Object lock = new Object();

  @GuardedBy("lock")
  private SpannerException exception;

  public MutationWriter(
      CopyTransactionMode transactionMode,
      Connection connection,
      String tableName,
      Map<String, Type> tableColumns,
      int indexedColumnsCount,
      CopyOptions.Format copyFormat,
      CSVFormat format,
      boolean hasHeader)
      throws IOException {
    this.transactionMode = transactionMode;
    this.connection = connection;
    this.tableName = tableName;
    this.tableColumns = tableColumns;
    int mutationLimit =
        Math.max(
            Integer.parseInt(
                System.getProperty(
                    "copy_in_mutation_limit", String.valueOf(DEFAULT_MUTATION_LIMIT))),
            1);
    this.maxBatchSize = Math.max(mutationLimit / (tableColumns.size() + indexedColumnsCount), 1);
    float commitLimitMultiplierFactor =
        Math.max(
            Float.parseFloat(
                System.getProperty(
                    "copy_in_commit_limit_multiplier_factor",
                    String.valueOf(DEFAULT_COMMIT_LIMIT_MULTIPLIER_FACTOR))),
            1.0f);
    this.commitSizeLimitForBatching =
        Math.round((float) commitSizeLimit / commitLimitMultiplierFactor);
    this.maxParallelism =
        Math.max(
            Integer.parseInt(
                System.getProperty(
                    "copy_in_max_parallelism", String.valueOf(DEFAULT_MAX_PARALLELISM))),
            1);
    int pipeBufferSize =
        Math.max(
            Integer.parseInt(
                System.getProperty(
                    "copy_in_pipe_buffer_size", String.valueOf(DEFAULT_PIPE_BUFFER_SIZE))),
            1024);
    this.parser = CopyInParser.create(copyFormat, format, payload, pipeBufferSize, hasHeader);
  }

  /** @return number of rows copied into Spanner */
  public long getRowCount() {
    return this.rowCount;
  }

  public void addCopyData(byte[] payload) {
    synchronized (lock) {
      if (this.exception != null) {
        throw this.exception;
      }
    }
    try {
      this.payload.write(payload);
    } catch (IOException e) {
      // Ignore the exception if the executor has already been shutdown. That means that an error
      // occurred that ended the COPY operation while we were writing data to the buffer.
      if (!executorService.isShutdown()) {
        SpannerException spannerException =
            SpannerExceptionFactory.newSpannerException(
                ErrorCode.INTERNAL, "Could not write copy data to buffer", e);
        logger.log(Level.SEVERE, spannerException.getMessage(), spannerException);
        throw spannerException;
      }
    }
  }

  /** Indicate that this mutation writer should commit. */
  public void commit() {
    this.commit.set(true);
  }

  /**
   * Indicate that this mutation writer should be rolled back. This will not rollback any changes
   * that have already been committed if the mutation writer is running in {@link
   * CopyTransactionMode#ImplicitNonAtomic}.
   */
  public void rollback() {
    this.rollback.set(true);
  }

  @Override
  public void close() throws IOException {
    this.payload.close();
    this.closedLatch.countDown();
  }

  @Override
  public StatementResult call() throws Exception {
    // This LinkedBlockingDeque holds a reference to all transactions that are currently active. The
    // max capacity of this deque is what ensures that we never have more than maxParallelism
    // transactions running at the same time. We could also achieve that by using a thread pool with
    // a fixed number of threads. The problem with that is however that Java does not have a thread
    // pool implementation that will block if a new task is offered and all threads are currently in
    // use. The only options are 'fail or add to queue'. We want to block our worker thread in this
    // case when the max parallel transactions has been reached, as that automatically creates back-
    // pressure in our entire pipeline that consists of:
    // Client app (psql) -> CopyData message -> CSVParser -> Transaction.
    LinkedBlockingDeque<ApiFuture<Void>> activeCommitFutures =
        new LinkedBlockingDeque<>(maxParallelism);
    // This list holds all transactions that we have started. We will wait on this entire list
    // before finishing, to ensure that all data has been written before we signal that we are done.
    List<ApiFuture<Void>> allCommitFutures = new ArrayList<>();
    try {
      Iterator<CopyRecord> iterator = this.parser.iterator();
      List<Mutation> mutations = new ArrayList<>();
      long currentBufferByteSize = 0L;
      // Note: iterator.hasNext() blocks if there is not enough data in the pipeline to construct a
      // complete record. It returns false if the stream has been closed and all records have been
      // returned.
      while (!rollback.get() && iterator.hasNext()) {
        CopyRecord record = iterator.next();
        if (record.numColumns() != this.tableColumns.keySet().size()) {
          throw SpannerExceptionFactory.newSpannerException(
              ErrorCode.INVALID_ARGUMENT,
              "Invalid COPY data: Row length mismatched. Expected "
                  + this.tableColumns.keySet().size()
                  + " columns, but only found "
                  + record.numColumns());
        }

        Mutation mutation = buildMutation(record);
        int mutationSize = calculateSize(mutation);
        this.rowCount++;

        if (transactionMode == CopyTransactionMode.ImplicitNonAtomic) {
          currentBufferByteSize =
              addMutationAndMaybeFlushTransaction(
                  activeCommitFutures,
                  allCommitFutures,
                  mutations,
                  mutation,
                  currentBufferByteSize,
                  mutationSize);
        } else {
          mutations.add(mutation);
          currentBufferByteSize += mutationSize;
          if (mutations.size() > maxBatchSize) {
            throw SpannerExceptionFactory.newSpannerException(
                ErrorCode.FAILED_PRECONDITION,
                "Record count: "
                    + mutations.size()
                    + " has exceeded the limit: "
                    + maxBatchSize
                    + ".\n\nThe number of mutations per record is equal to the number of columns in the record "
                    + "plus the number of indexed columns in the record. The maximum number of mutations "
                    + "in one transaction is "
                    + DEFAULT_MUTATION_LIMIT
                    + ".\n\nExecute `SET SPANNER.AUTOCOMMIT_DML_MODE='PARTITIONED_NON_ATOMIC'` before executing a large COPY operation "
                    + "to instruct PGAdapter to automatically break large transactions into multiple smaller. "
                    + "This will make the COPY operation non-atomic.\n\n");
          }
          if (currentBufferByteSize > commitSizeLimit) {
            throw SpannerExceptionFactory.newSpannerException(
                ErrorCode.FAILED_PRECONDITION,
                "Commit size: "
                    + currentBufferByteSize
                    + " has exceeded the limit: "
                    + commitSizeLimit
                    + ".\n\nExecute `SET SPANNER.AUTOCOMMIT_DML_MODE='PARTITIONED_NON_ATOMIC'` before executing a large COPY operation "
                    + "to instruct PGAdapter to automatically break large transactions into multiple smaller. "
                    + "This will make the COPY operation non-atomic.\n\n");
          }
        }
      } // end of iterator.hasNext()

      // There are no more CSVRecords in the pipeline.
      // Write any remaining mutations in the buffer.
      if (!rollback.get() && !mutations.isEmpty()) {
        if (transactionMode == CopyTransactionMode.Explicit) {
          connection.write(mutations);
        } else {
          // Wait until we have received a CopyDone message before writing the remaining data to
          // Spanner. If we are in a non-atomic transaction, there might already be data that have
          // been written to Spanner.
          closedLatch.await();
          if (commit.get()) {
            allCommitFutures.add(writeToSpannerAsync(activeCommitFutures, mutations));
          }
        }
      }
      // Wait for all commits to finish. We do this even if something went wrong, as it ensures two
      // things:
      // 1. All commits that were in flight when something went wrong will finish before we return
      //    an error to the client application. This prevents commits still being added to the
      //    database after we have returned an error, which could cause confusion.
      // 2. This will throw the underlying exception, so we can catch and register it.
      ApiFutures.allAsList(allCommitFutures).get();
    } catch (SpannerException e) {
      synchronized (lock) {
        this.exception = e;
        throw this.exception;
      }
    } catch (ExecutionException e) {
      synchronized (lock) {
        this.exception = SpannerExceptionFactory.asSpannerException(e.getCause());
        throw this.exception;
      }
    } catch (Exception e) {
      synchronized (lock) {
        this.exception = SpannerExceptionFactory.asSpannerException(e);
        throw this.exception;
      }
    } finally {
      this.executorService.shutdown();
      if (!this.executorService.awaitTermination(60L, TimeUnit.SECONDS)) {
        logger.log(Level.WARNING, "Timeout while waiting for MutationWriter executor to shutdown.");
      }
      this.payload.close();
      this.parser.close();
    }
    return new UpdateCount(rowCount);
  }

  private long addMutationAndMaybeFlushTransaction(
      LinkedBlockingDeque<ApiFuture<Void>> activeCommitFutures,
      List<ApiFuture<Void>> allCommitFutures,
      List<Mutation> mutations,
      Mutation mutation,
      long currentBufferByteSize,
      int mutationSize)
      throws Exception {
    // Flush before adding to the buffer if the commit size would be exceeded.
    if (!mutations.isEmpty() && currentBufferByteSize + mutationSize > commitSizeLimitForBatching) {
      allCommitFutures.add(writeToSpannerAsync(activeCommitFutures, mutations));
      mutations.clear();
      mutations.add(mutation);
      return mutationSize;
    }

    mutations.add(mutation);
    if (mutations.size() == maxBatchSize) {
      allCommitFutures.add(writeToSpannerAsync(activeCommitFutures, mutations));
      mutations.clear();
      return 0L; // Buffer is empty, so the batch size in bytes is now back to zero.
    }
    return currentBufferByteSize + mutationSize;
  }

  private ApiFuture<Void> writeToSpannerAsync(
      LinkedBlockingDeque<ApiFuture<Void>> activeCommitFutures, Iterable<Mutation> mutations)
      throws Exception {

    SettableApiFuture<Void> settableApiFuture = SettableApiFuture.create();
    // Add this future to the list of active commit futures. This will block if the deque is full,
    // and this will effectively apply back-pressure to the entire stream as the worker thread is
    // blocked until there is room in the deque.
    activeCommitFutures.put(settableApiFuture);

    DatabaseClient dbClient = connection.getDatabaseClient();
    ImmutableList<Mutation> immutableMutations = ImmutableList.copyOf(mutations);
    ListenableFuture<Void> listenableFuture =
        executorService.submit(
            () -> {
              Context context =
                  Context.current()
                      .withValue(
                          SpannerOptions.CALL_CONTEXT_CONFIGURATOR_KEY,
                          SpannerCallContextTimeoutConfigurator.create()
                              .withCommitTimeout(Duration.ofSeconds(100L)));
              context.run(() -> dbClient.write(immutableMutations));
              return null;
            });
    Futures.addCallback(
        listenableFuture,
        new FutureCallback<Void>() {
          @Override
          public void onFailure(@Nonnull Throwable t) {
            rollback.set(true);
            //noinspection ResultOfMethodCallIgnored
            activeCommitFutures.remove(settableApiFuture);
            settableApiFuture.setException(t);
          }

          @Override
          public void onSuccess(Void result) {
            //noinspection ResultOfMethodCallIgnored
            activeCommitFutures.remove(settableApiFuture);
{'fraction_non_alphanumeric': 0.07748331533327794, 'fraction_numerical': 0.003073855611863421, 'mean_word_length': 2.850303870348651, 'pattern_counts': {'":': 0, '<': 57, '<?xml version=': 0, '>': 68, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4335079', 'n_tokens_mistral': 9453, 'n_tokens_neox': 9295, 'n_words': 2890}