From 3c05037e8fce8d3ce4907fde85bd254fc98ea588 Mon Sep 17 00:00:00 2001
From: Brian Chen <chenbrian@google.com>
Date: Wed, 13 Jan 2021 18:43:03 -0600
Subject: [PATCH] feat: add success and error callbacks to BulkWriter (#483)

---
 .../clirr-ignored-differences.xml             |  19 +
 .../cloud/firestore/BulkCommitBatch.java      | 172 +++-
 .../google/cloud/firestore/BulkWriter.java    | 911 +++++++++++++-----
 .../cloud/firestore/BulkWriterException.java  |  68 ++
 .../com/google/cloud/firestore/Firestore.java |   3 +
 .../cloud/firestore/FirestoreException.java   |   4 +-
 .../google/cloud/firestore/FirestoreImpl.java |   8 +-
 .../google/cloud/firestore/Transaction.java   |   2 +-
 .../google/cloud/firestore/UpdateBuilder.java | 228 +----
 .../google/cloud/firestore/WriteBatch.java    |   2 +-
 .../cloud/firestore/BulkWriterTest.java       | 504 +++++++++-
 .../cloud/firestore/ITBulkWriterTest.java     | 216 +++++
 .../cloud/firestore/WriteBatchTest.java       |  41 -
 13 files changed, 1619 insertions(+), 559 deletions(-)
 create mode 100644 google-cloud-firestore/src/main/java/com/google/cloud/firestore/BulkWriterException.java
 create mode 100644 google-cloud-firestore/src/test/java/com/google/cloud/firestore/ITBulkWriterTest.java

diff --git a/google-cloud-firestore/clirr-ignored-differences.xml b/google-cloud-firestore/clirr-ignored-differences.xml
index a66036b79..4a4d5d4c0 100644
--- a/google-cloud-firestore/clirr-ignored-differences.xml
+++ b/google-cloud-firestore/clirr-ignored-differences.xml
@@ -223,6 +223,25 @@
     <to>java.util.List</to>
   </difference>
 
+  <!--
+  UpdateBuilder
+  -->
+  <difference>
+    <differenceType>6001</differenceType>
+    <className>com/google/cloud/firestore/UpdateBuilder</className>
+    <field>pendingOperations</field>
+  </difference>
+  <difference>
+    <differenceType>6001</differenceType>
+    <className>com/google/cloud/firestore/UpdateBuilder</className>
+    <field>state</field>
+  </difference>
+  <difference>
+    <differenceType>6010</differenceType>
+    <className>com/google/cloud/firestore/UpdateBuilder</className>
+    <field>writes</field>
+  </difference>
+
   <!--
   FakeCredentials Refactor
   com.google.cloud.firestore.FirestoreOptions$Builder$FakeCredentials -> com.google.cloud.firestore.FirestoreOptions$EmulatorCredentials
diff --git a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/BulkCommitBatch.java b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/BulkCommitBatch.java
index 8ef106fc3..af2180cee 100644
--- a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/BulkCommitBatch.java
+++ b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/BulkCommitBatch.java
@@ -16,32 +16,180 @@
 
 package com.google.cloud.firestore;
 
+import com.google.api.core.ApiAsyncFunction;
+import com.google.api.core.ApiFunction;
 import com.google.api.core.ApiFuture;
+import com.google.api.core.ApiFutures;
+import com.google.api.core.SettableApiFuture;
+import com.google.cloud.Timestamp;
 import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.util.concurrent.MoreExecutors;
+import com.google.firestore.v1.BatchWriteRequest;
+import com.google.firestore.v1.BatchWriteResponse;
+import io.grpc.Status;
+import io.opencensus.trace.AttributeValue;
+import io.opencensus.trace.Tracing;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+import java.util.concurrent.CopyOnWriteArraySet;
+import javax.annotation.Nullable;
 
 /** Used to represent a batch on the BatchQueue. */
 class BulkCommitBatch extends UpdateBuilder<ApiFuture<WriteResult>> {
+  /**
+   * Used to represent the state of batch.
+   *
+   * <p>Writes can only be added while the batch is OPEN. For a batch to be sent, the batch must be
+   * READY_TO_SEND. After a batch is sent, it is marked as SENT.
+   */
+  enum BatchState {
+    OPEN,
+    READY_TO_SEND,
+    SENT,
+  }
+
+  private BatchState state = BatchState.OPEN;
+
+  private final List<SettableApiFuture<BatchWriteResult>> pendingOperations = new ArrayList<>();
+  private final Set<DocumentReference> documents = new CopyOnWriteArraySet<>();
+  private final int maxBatchSize;
 
   BulkCommitBatch(FirestoreImpl firestore, int maxBatchSize) {
-    super(firestore, maxBatchSize);
+    super(firestore);
+    this.maxBatchSize = maxBatchSize;
   }
 
-  BulkCommitBatch(FirestoreImpl firestore, BulkCommitBatch retryBatch) {
-    super(firestore);
+  @Override
+  boolean isCommitted() {
+    return state == BatchState.SENT;
+  }
+
+  ApiFuture<WriteResult> wrapResult(DocumentReference documentReference) {
+    return processLastOperation(documentReference);
+  }
+
+  /**
+   * Commits all pending operations to the database and verifies all preconditions.
+   *
+   * <p>The writes in the batch are not applied atomically and can be applied out of order.
+   */
+  ApiFuture<List<BatchWriteResult>> bulkCommit() {
+    Tracing.getTracer()
+        .getCurrentSpan()
+        .addAnnotation(
+            TraceUtil.SPAN_NAME_BATCHWRITE,
+            ImmutableMap.of("numDocuments", AttributeValue.longAttributeValue(getWrites().size())));
 
-    // Create a new BulkCommitBatch containing only the indexes from the provided indexes to retry.
-    for (int index : retryBatch.getPendingIndexes()) {
-      this.writes.add(retryBatch.writes.get(index));
+    Preconditions.checkState(
+        isReadyToSend(), "The batch should be marked as READY_TO_SEND before committing");
+    state = BatchState.SENT;
+
+    final BatchWriteRequest.Builder request = BatchWriteRequest.newBuilder();
+    request.setDatabase(firestore.getDatabaseName());
+
+    for (WriteOperation writeOperation : getWrites()) {
+      request.addWrites(writeOperation.write);
     }
 
+    ApiFuture<BatchWriteResponse> response =
+        firestore.sendRequest(request.build(), firestore.getClient().batchWriteCallable());
+
+    return ApiFutures.transform(
+        response,
+        new ApiFunction<BatchWriteResponse, List<BatchWriteResult>>() {
+          @Override
+          public List<BatchWriteResult> apply(BatchWriteResponse batchWriteResponse) {
+            List<com.google.firestore.v1.WriteResult> writeResults =
+                batchWriteResponse.getWriteResultsList();
+
+            List<com.google.rpc.Status> statuses = batchWriteResponse.getStatusList();
+
+            List<BatchWriteResult> result = new ArrayList<>();
+
+            for (int i = 0; i < writeResults.size(); ++i) {
+              com.google.firestore.v1.WriteResult writeResult = writeResults.get(i);
+              com.google.rpc.Status status = statuses.get(i);
+              Status code = Status.fromCodeValue(status.getCode());
+              @Nullable Timestamp updateTime = null;
+              @Nullable Exception exception = null;
+              if (code == Status.OK) {
+                updateTime = Timestamp.fromProto(writeResult.getUpdateTime());
+              } else {
+                exception = FirestoreException.serverRejected(code, status.getMessage());
+              }
+              result.add(new BatchWriteResult(updateTime, exception));
+            }
+
+            return result;
+          }
+        },
+        MoreExecutors.directExecutor());
+  }
+
+  int getPendingOperationCount() {
+    return pendingOperations.size();
+  }
+
+  ApiFuture<WriteResult> processLastOperation(DocumentReference documentReference) {
     Preconditions.checkState(
-        retryBatch.state == BatchState.SENT,
-        "Batch should be SENT when creating a new BulkCommitBatch for retry");
-    this.state = retryBatch.state;
-    this.pendingOperations = retryBatch.pendingOperations;
+        !documents.contains(documentReference),
+        "Batch should not contain writes to the same document");
+    documents.add(documentReference);
+    Preconditions.checkState(state == BatchState.OPEN, "Batch should be OPEN when adding writes");
+    SettableApiFuture<BatchWriteResult> resultFuture = SettableApiFuture.create();
+    pendingOperations.add(resultFuture);
+
+    if (getPendingOperationCount() == maxBatchSize) {
+      state = BatchState.READY_TO_SEND;
+    }
+
+    return ApiFutures.transformAsync(
+        resultFuture,
+        new ApiAsyncFunction<BatchWriteResult, WriteResult>() {
+          public ApiFuture<WriteResult> apply(BatchWriteResult batchWriteResult) throws Exception {
+            if (batchWriteResult.getException() == null) {
+              return ApiFutures.immediateFuture(new WriteResult(batchWriteResult.getWriteTime()));
+            } else {
+              throw batchWriteResult.getException();
+            }
+          }
+        },
+        MoreExecutors.directExecutor());
+  }
+
+  /**
+   * Resolves the individual operations in the batch with the results and removes the entry from the
+   * pendingOperations map if the result is not retryable.
+   */
+  void processResults(List<BatchWriteResult> results) {
+    for (int i = 0; i < results.size(); i++) {
+      SettableApiFuture<BatchWriteResult> resultFuture = pendingOperations.get(i);
+      BatchWriteResult result = results.get(i);
+      if (result.getException() == null) {
+        resultFuture.set(result);
+      } else {
+        resultFuture.setException(result.getException());
+      }
+    }
+  }
+
+  void markReadyToSend() {
+    if (state == BatchState.OPEN) {
+      state = BatchState.READY_TO_SEND;
+    }
+  }
+
+  boolean isOpen() {
+    return state == BatchState.OPEN;
+  }
+
+  boolean isReadyToSend() {
+    return state == BatchState.READY_TO_SEND;
   }
 
-  ApiFuture<WriteResult> wrapResult(ApiFuture<WriteResult> result) {
-    return result;
+  boolean has(DocumentReference documentReference) {
+    return documents.contains(documentReference);
   }
 }
diff --git a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/BulkWriter.java b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/BulkWriter.java
index 8fad7dea2..6dd86c720 100644
--- a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/BulkWriter.java
+++ b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/BulkWriter.java
@@ -19,21 +19,21 @@
 import com.google.api.core.ApiAsyncFunction;
 import com.google.api.core.ApiFuture;
 import com.google.api.core.ApiFutures;
-import com.google.api.core.CurrentMillisClock;
 import com.google.api.core.SettableApiFuture;
-import com.google.api.gax.retrying.ExponentialRetryAlgorithm;
-import com.google.api.gax.retrying.TimedAttemptSettings;
-import com.google.cloud.firestore.UpdateBuilder.BatchState;
+import com.google.api.gax.rpc.StatusCode.Code;
+import com.google.cloud.firestore.v1.FirestoreSettings;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
-import com.google.common.base.Predicate;
-import com.google.common.collect.FluentIterable;
+import com.google.common.collect.Lists;
 import com.google.common.util.concurrent.MoreExecutors;
 import java.util.ArrayList;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
-import java.util.concurrent.CopyOnWriteArrayList;
+import java.util.Set;
 import java.util.concurrent.ExecutionException;
+import java.util.concurrent.Executor;
+import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
 import java.util.logging.Level;
@@ -42,9 +42,47 @@
 import javax.annotation.Nullable;
 
 final class BulkWriter implements AutoCloseable {
+  /**
+   * A callback set by `addWriteResultListener()` to be run every time an operation successfully
+   * completes.
+   */
+  public interface WriteResultCallback {
+    /**
+     * @param documentReference The document the operation was performed on.
+     * @param result The result of the operation.
+     */
+    void onResult(DocumentReference documentReference, WriteResult result);
+  };
+
+  /**
+   * A callback set by `addWriteErrorListener()` to be run every time an operation fails and
+   * determines if an operation should be retried.
+   */
+  public interface WriteErrorCallback {
+    /**
+     * @param error The error object from the failed BulkWriter operation attempt.
+     * @return Whether to retry the operation or not.
+     */
+    boolean onError(BulkWriterException error);
+  }
+
+  private interface BulkWriterOperationCallback {
+    ApiFuture<WriteResult> apply(BulkCommitBatch batch);
+  }
+
+  enum OperationType {
+    CREATE,
+    SET,
+    UPDATE,
+    DELETE
+  }
   /** The maximum number of writes that can be in a single batch. */
-  public static final int MAX_BATCH_SIZE = 500;
+  public static final int MAX_BATCH_SIZE = 20;
 
+  /**
+   * The maximum number of retries that will be attempted with backoff before stopping all retry
+   * attempts.
+   */
   public static final int MAX_RETRY_ATTEMPTS = 10;
 
   /**
@@ -72,17 +110,51 @@ final class BulkWriter implements AutoCloseable {
    */
   private static final int RATE_LIMITER_MULTIPLIER_MILLIS = 5 * 60 * 1000;
 
+  private static final WriteResultCallback DEFAULT_SUCCESS_LISTENER =
+      new WriteResultCallback() {
+        public void onResult(DocumentReference documentReference, WriteResult result) {}
+      };
+
+  private static final WriteErrorCallback DEFAULT_ERROR_LISTENER =
+      new WriteErrorCallback() {
+        public boolean onError(BulkWriterException error) {
+          if (error.getFailedAttempts() == MAX_RETRY_ATTEMPTS) {
+            return false;
+          }
+          Set<Code> codes = FirestoreSettings.newBuilder().batchWriteSettings().getRetryableCodes();
+          for (Code code : codes) {
+            if (code.equals(Code.valueOf(error.getStatus().getCode().name()))) {
+              return true;
+            }
+          }
+          return false;
+        }
+      };
+
   private static final Logger logger = Logger.getLogger(BulkWriter.class.getName());
 
   /** The maximum number of writes that can be in a single batch. */
   private int maxBatchSize = MAX_BATCH_SIZE;
 
+  /** A queue of batches to be written. */
+  private final List<BulkCommitBatch> batchQueue = new ArrayList<>();
+
+  /** A queue of batches containing operations that need to be retried. */
+  private final List<BulkCommitBatch> retryBatchQueue = new ArrayList<>();
+
+  /**
+   * A set of futures that represent pending BulkWriter operations. Each future is completed when
+   * the BulkWriter operation resolves with either a WriteResult or failure after all applicable
+   * retries are performed. This set of futures is used by flush() to track when enqueued writes are
+   * completed.
+   */
+  private final Set<ApiFuture<Void>> pendingOperations = new HashSet<>();
+
   /**
-   * A queue of batches to be written. Use a synchronized list to avoid multi-thread concurrent
-   * modification errors (as this list is modified from both the user thread and the network
-   * thread).
+   * A list of futures that represent sent batches. Each future is completed when the batch's
+   * response is received. This includes batches from both the batchQueue and retryBatchQueue.
    */
-  private final List<BulkCommitBatch> batchQueue = new CopyOnWriteArrayList<>();
+  private final Set<ApiFuture<Void>> pendingBatches = new HashSet<>();
 
   /** Whether this BulkWriter instance is closed. Once closed, it cannot be opened again. */
   private boolean closed = false;
@@ -90,20 +162,37 @@ final class BulkWriter implements AutoCloseable {
   /** Rate limiter used to throttle requests as per the 500/50/5 rule. */
   private final RateLimiter rateLimiter;
 
-  private final FirestoreImpl firestore;
+  private WriteResultCallback successListener = DEFAULT_SUCCESS_LISTENER;
+
+  private WriteErrorCallback errorListener = DEFAULT_ERROR_LISTENER;
+
+  private Executor successExecutor;
+  private Executor errorExecutor;
 
-  private final ScheduledExecutorService firestoreExecutor;
+  /**
+   * Used to track when writes are enqueued. The user handler executors cannot be changed after a
+   * write has been enqueued.
+   */
+  private boolean writesEnqueued = false;
 
-  private final ExponentialRetryAlgorithm backoff;
-  private TimedAttemptSettings nextAttempt;
+  private final FirestoreImpl firestore;
+
+  // Executor used to run all BulkWriter operations. BulkWriter uses its own executor since we
+  // don't want to block a gax/grpc executor while running user error and success callbacks.
+  private final ScheduledExecutorService bulkWriterExecutor;
 
   BulkWriter(FirestoreImpl firestore, BulkWriterOptions options) {
+    this(firestore, options, Executors.newSingleThreadScheduledExecutor());
+  }
+
+  BulkWriter(
+      FirestoreImpl firestore,
+      BulkWriterOptions options,
+      ScheduledExecutorService bulkWriterExecutor) {
     this.firestore = firestore;
-    this.backoff =
-        new ExponentialRetryAlgorithm(
-            firestore.getOptions().getRetrySettings(), CurrentMillisClock.getDefaultClock());
-    this.nextAttempt = backoff.createFirstAttempt();
-    this.firestoreExecutor = firestore.getClient().getExecutor();
+    this.bulkWriterExecutor = bulkWriterExecutor;
+    this.successExecutor = MoreExecutors.directExecutor();
+    this.errorExecutor = MoreExecutors.directExecutor();
 
     if (!options.getThrottlingEnabled()) {
       this.rateLimiter =
@@ -149,16 +238,22 @@ final class BulkWriter implements AutoCloseable {
    *
    * @param documentReference A reference to the document to be created.
    * @param fields A map of the fields and values for the document.
-   * @return An ApiFuture containing the result of the write. Contains an error if the write fails.
+   * @return An ApiFuture containing the result of the write. Contains a {@link BulkWriterException}
+   *     if the write fails.
    */
   @Nonnull
   public ApiFuture<WriteResult> create(
-      @Nonnull DocumentReference documentReference, @Nonnull Map<String, Object> fields) {
+      @Nonnull final DocumentReference documentReference,
+      @Nonnull final Map<String, Object> fields) {
     verifyNotClosed();
-    BulkCommitBatch bulkCommitBatch = getEligibleBatch();
-    ApiFuture<WriteResult> future = bulkCommitBatch.create(documentReference, fields);
-    sendReadyBatches();
-    return future;
+    return executeWrite(
+        documentReference,
+        OperationType.CREATE,
+        new BulkWriterOperationCallback() {
+          public ApiFuture<WriteResult> apply(BulkCommitBatch batch) {
+            return batch.create(documentReference, fields);
+          }
+        });
   }
 
   /**
@@ -167,32 +262,41 @@ public ApiFuture<WriteResult> create(
    *
    * @param documentReference A reference to the document to be created.
    * @param pojo The POJO that will be used to populate the document contents.
-   * @return An ApiFuture containing the result of the write. Contains an error if the write fails.
+   * @return An ApiFuture containing the result of the write. Contains a {@link BulkWriterException}
+   *     if the write fails.
    */
   @Nonnull
   public ApiFuture<WriteResult> create(
-      @Nonnull DocumentReference documentReference, @Nonnull Object pojo) {
+      @Nonnull final DocumentReference documentReference, @Nonnull final Object pojo) {
     verifyNotClosed();
-    BulkCommitBatch bulkCommitBatch = getEligibleBatch();
-    ApiFuture<WriteResult> future = bulkCommitBatch.create(documentReference, pojo);
-    sendReadyBatches();
-    return future;
+    return executeWrite(
+        documentReference,
+        OperationType.CREATE,
+        new BulkWriterOperationCallback() {
+          public ApiFuture<WriteResult> apply(BulkCommitBatch batch) {
+            return batch.create(documentReference, pojo);
+          }
+        });
   }
 
   /**
    * Delete a document from the database.
    *
    * @param documentReference The DocumentReference to delete.
-   * @return An ApiFuture containing the result of the delete. Contains an error if the delete
-   *     fails.
+   * @return An ApiFuture containing a sentinel value (Timestamp(0)) for the delete operation.
+   *     Contains a {@link BulkWriterException} if the delete fails.
    */
   @Nonnull
-  public ApiFuture<WriteResult> delete(@Nonnull DocumentReference documentReference) {
+  public ApiFuture<WriteResult> delete(@Nonnull final DocumentReference documentReference) {
     verifyNotClosed();
-    BulkCommitBatch bulkCommitBatch = getEligibleBatch();
-    ApiFuture<WriteResult> future = bulkCommitBatch.delete(documentReference);
-    sendReadyBatches();
-    return future;
+    return executeWrite(
+        documentReference,
+        OperationType.DELETE,
+        new BulkWriterOperationCallback() {
+          public ApiFuture<WriteResult> apply(BulkCommitBatch batch) {
+            return batch.delete(documentReference);
+          }
+        });
   }
 
   /**
@@ -201,16 +305,21 @@ public ApiFuture<WriteResult> delete(@Nonnull DocumentReference documentReferenc
    * @param documentReference The DocumentReference to delete.
    * @param precondition Precondition to enforce for this delete.
    * @return An ApiFuture containing a sentinel value (Timestamp(0)) for the delete operation.
-   *     Contains an error if the delete fails.
+   *     Contains a {@link BulkWriterException} if the delete fails.
    */
   @Nonnull
   public ApiFuture<WriteResult> delete(
-      @Nonnull DocumentReference documentReference, @Nonnull Precondition precondition) {
+      @Nonnull final DocumentReference documentReference,
+      @Nonnull final Precondition precondition) {
     verifyNotClosed();
-    BulkCommitBatch bulkCommitBatch = getEligibleBatch();
-    ApiFuture<WriteResult> future = bulkCommitBatch.delete(documentReference, precondition);
-    sendReadyBatches();
-    return future;
+    return executeWrite(
+        documentReference,
+        OperationType.DELETE,
+        new BulkWriterOperationCallback() {
+          public ApiFuture<WriteResult> apply(BulkCommitBatch batch) {
+            return batch.delete(documentReference, precondition);
+          }
+        });
   }
 
   /**
@@ -219,17 +328,22 @@ public ApiFuture<WriteResult> delete(
    *
    * @param documentReference A reference to the document to be set.
    * @param fields A map of the fields and values for the document.
-   * @return An ApiFuture containing a sentinel value (Timestamp(0)) for the delete operation.
-   *     Contains an error if the delete fails.
+   * @return An ApiFuture containing the result of the write. Contains a {@link BulkWriterException}
+   *     if the write fails.
    */
   @Nonnull
   public ApiFuture<WriteResult> set(
-      @Nonnull DocumentReference documentReference, @Nonnull Map<String, Object> fields) {
+      @Nonnull final DocumentReference documentReference,
+      @Nonnull final Map<String, Object> fields) {
     verifyNotClosed();
-    BulkCommitBatch bulkCommitBatch = getEligibleBatch();
-    ApiFuture<WriteResult> future = bulkCommitBatch.set(documentReference, fields);
-    sendReadyBatches();
-    return future;
+    return executeWrite(
+        documentReference,
+        OperationType.SET,
+        new BulkWriterOperationCallback() {
+          public ApiFuture<WriteResult> apply(BulkCommitBatch batch) {
+            return batch.set(documentReference, fields);
+          }
+        });
   }
 
   /**
@@ -240,18 +354,23 @@ public ApiFuture<WriteResult> set(
    * @param documentReference A reference to the document to be set.
    * @param fields A map of the fields and values for the document.
    * @param options An object to configure the set behavior.
-   * @return An ApiFuture containing the result of the write. Contains an error if the write fails.
+   * @return An ApiFuture containing the result of the write. Contains a {@link BulkWriterException}
+   *     if the write fails.
    */
   @Nonnull
   public ApiFuture<WriteResult> set(
-      @Nonnull DocumentReference documentReference,
-      @Nonnull Map<String, Object> fields,
-      @Nonnull SetOptions options) {
+      @Nonnull final DocumentReference documentReference,
+      @Nonnull final Map<String, Object> fields,
+      @Nonnull final SetOptions options) {
     verifyNotClosed();
-    BulkCommitBatch bulkCommitBatch = getEligibleBatch();
-    ApiFuture<WriteResult> future = bulkCommitBatch.set(documentReference, fields, options);
-    sendReadyBatches();
-    return future;
+    return executeWrite(
+        documentReference,
+        OperationType.SET,
+        new BulkWriterOperationCallback() {
+          public ApiFuture<WriteResult> apply(BulkCommitBatch batch) {
+            return batch.set(documentReference, fields, options);
+          }
+        });
   }
 
   /**
@@ -262,16 +381,23 @@ public ApiFuture<WriteResult> set(
    * @param documentReference A reference to the document to be set.
    * @param pojo The POJO that will be used to populate the document contents.
    * @param options An object to configure the set behavior.
-   * @return An ApiFuture containing the result of the write. Contains an error if the write fails.
+   * @return An ApiFuture containing the result of the write. Contains a {@link BulkWriterException}
+   *     if the write fails.
    */
   @Nonnull
   public ApiFuture<WriteResult> set(
-      @Nonnull DocumentReference documentReference, Object pojo, @Nonnull SetOptions options) {
+      @Nonnull final DocumentReference documentReference,
+      @Nonnull final Object pojo,
+      @Nonnull final SetOptions options) {
     verifyNotClosed();
-    BulkCommitBatch bulkCommitBatch = getEligibleBatch();
-    ApiFuture<WriteResult> future = bulkCommitBatch.set(documentReference, pojo, options);
-    sendReadyBatches();
-    return future;
+    return executeWrite(
+        documentReference,
+        OperationType.SET,
+        new BulkWriterOperationCallback() {
+          public ApiFuture<WriteResult> apply(BulkCommitBatch batch) {
+            return batch.set(documentReference, pojo, options);
+          }
+        });
   }
 
   /**
@@ -280,15 +406,21 @@ public ApiFuture<WriteResult> set(
    *
    * @param documentReference A reference to the document to be set.
    * @param pojo The POJO that will be used to populate the document contents.
-   * @return An ApiFuture containing the result of the write. Contains an error if the write fails.
+   * @return An ApiFuture containing the result of the write. Contains a {@link BulkWriterException}
+   *     if the write fails.
    */
   @Nonnull
-  public ApiFuture<WriteResult> set(@Nonnull DocumentReference documentReference, Object pojo) {
+  public ApiFuture<WriteResult> set(
+      @Nonnull final DocumentReference documentReference, @Nonnull final Object pojo) {
     verifyNotClosed();
-    BulkCommitBatch bulkCommitBatch = getEligibleBatch();
-    ApiFuture<WriteResult> future = bulkCommitBatch.set(documentReference, pojo);
-    sendReadyBatches();
-    return future;
+    return executeWrite(
+        documentReference,
+        OperationType.SET,
+        new BulkWriterOperationCallback() {
+          public ApiFuture<WriteResult> apply(BulkCommitBatch batch) {
+            return batch.set(documentReference, pojo);
+          }
+        });
   }
 
   /**
@@ -302,16 +434,22 @@ public ApiFuture<WriteResult> set(@Nonnull DocumentReference documentReference,
    *
    * @param documentReference A reference to the document to be updated.
    * @param fields A map of the fields and values for the document.
-   * @return An ApiFuture containing the result of the write. Contains an error if the write fails.
+   * @return An ApiFuture containing the result of the write. Contains a {@link BulkWriterException}
+   *     if the write fails.
    */
   @Nonnull
   public ApiFuture<WriteResult> update(
-      @Nonnull DocumentReference documentReference, @Nonnull Map<String, Object> fields) {
+      @Nonnull final DocumentReference documentReference,
+      @Nonnull final Map<String, Object> fields) {
     verifyNotClosed();
-    BulkCommitBatch bulkCommitBatch = getEligibleBatch();
-    ApiFuture<WriteResult> future = bulkCommitBatch.update(documentReference, fields);
-    sendReadyBatches();
-    return future;
+    return executeWrite(
+        documentReference,
+        OperationType.UPDATE,
+        new BulkWriterOperationCallback() {
+          public ApiFuture<WriteResult> apply(BulkCommitBatch batch) {
+            return batch.update(documentReference, fields);
+          }
+        });
   }
 
   /**
@@ -326,18 +464,23 @@ public ApiFuture<WriteResult> update(
    * @param documentReference A reference to the document to be updated.
    * @param fields A map of the fields and values for the document.
    * @param precondition Precondition to enforce on this update.
-   * @return An ApiFuture containing the result of the write. Contains an error if the write fails.
+   * @return An ApiFuture containing the result of the write. Contains a {@link BulkWriterException}
+   *     if the write fails.
    */
   @Nonnull
   public ApiFuture<WriteResult> update(
-      @Nonnull DocumentReference documentReference,
-      @Nonnull Map<String, Object> fields,
-      Precondition precondition) {
+      @Nonnull final DocumentReference documentReference,
+      @Nonnull final Map<String, Object> fields,
+      @Nonnull final Precondition precondition) {
     verifyNotClosed();
-    BulkCommitBatch bulkCommitBatch = getEligibleBatch();
-    ApiFuture<WriteResult> future = bulkCommitBatch.update(documentReference, fields, precondition);
-    sendReadyBatches();
-    return future;
+    return executeWrite(
+        documentReference,
+        OperationType.UPDATE,
+        new BulkWriterOperationCallback() {
+          public ApiFuture<WriteResult> apply(BulkCommitBatch batch) {
+            return batch.update(documentReference, fields, precondition);
+          }
+        });
   }
 
   /**
@@ -353,20 +496,24 @@ public ApiFuture<WriteResult> update(
    * @param field The first field to set.
    * @param value The first value to set.
    * @param moreFieldsAndValues String and Object pairs with more fields to be set.
-   * @return An ApiFuture containing the result of the write. Contains an error if the write fails.
+   * @return An ApiFuture containing the result of the write. Contains a {@link BulkWriterException}
+   *     if the write fails.
    */
   @Nonnull
   public ApiFuture<WriteResult> update(
-      @Nonnull DocumentReference documentReference,
-      @Nonnull String field,
-      @Nullable Object value,
-      Object... moreFieldsAndValues) {
+      @Nonnull final DocumentReference documentReference,
+      @Nonnull final String field,
+      @Nullable final Object value,
+      final Object... moreFieldsAndValues) {
     verifyNotClosed();
-    BulkCommitBatch bulkCommitBatch = getEligibleBatch();
-    ApiFuture<WriteResult> future =
-        bulkCommitBatch.update(documentReference, field, value, moreFieldsAndValues);
-    sendReadyBatches();
-    return future;
+    return executeWrite(
+        documentReference,
+        OperationType.UPDATE,
+        new BulkWriterOperationCallback() {
+          public ApiFuture<WriteResult> apply(BulkCommitBatch batch) {
+            return batch.update(documentReference, field, value, moreFieldsAndValues);
+          }
+        });
   }
 
   /**
@@ -382,20 +529,24 @@ public ApiFuture<WriteResult> update(
    * @param fieldPath The first field to set.
    * @param value The first value to set.
    * @param moreFieldsAndValues String and Object pairs with more fields to be set.
-   * @return An ApiFuture containing the result of the write. Contains an error if the write fails.
+   * @return An ApiFuture containing the result of the write. Contains a {@link BulkWriterException}
+   *     if the write fails.
    */
   @Nonnull
   public ApiFuture<WriteResult> update(
-      @Nonnull DocumentReference documentReference,
-      @Nonnull FieldPath fieldPath,
-      @Nullable Object value,
-      Object... moreFieldsAndValues) {
+      @Nonnull final DocumentReference documentReference,
+      @Nonnull final FieldPath fieldPath,
+      @Nullable final Object value,
+      final Object... moreFieldsAndValues) {
     verifyNotClosed();
-    BulkCommitBatch bulkCommitBatch = getEligibleBatch();
-    ApiFuture<WriteResult> future =
-        bulkCommitBatch.update(documentReference, fieldPath, value, moreFieldsAndValues);
-    sendReadyBatches();
-    return future;
+    return executeWrite(
+        documentReference,
+        OperationType.UPDATE,
+        new BulkWriterOperationCallback() {
+          public ApiFuture<WriteResult> apply(BulkCommitBatch batch) {
+            return batch.update(documentReference, fieldPath, value, moreFieldsAndValues);
+          }
+        });
   }
 
   /**
@@ -411,21 +562,25 @@ public ApiFuture<WriteResult> update(
    * @param field The first field to set.
    * @param value The first value to set.
    * @param moreFieldsAndValues String and Object pairs with more fields to be set.
-   * @return An ApiFuture containing the result of the write. Contains an error if the write fails.
+   * @return An ApiFuture containing the result of the write. Contains a {@link BulkWriterException}
+   *     if the write fails.
    */
   @Nonnull
   public ApiFuture<WriteResult> update(
-      @Nonnull DocumentReference documentReference,
-      @Nonnull Precondition precondition,
-      @Nonnull String field,
-      @Nullable Object value,
-      Object... moreFieldsAndValues) {
+      @Nonnull final DocumentReference documentReference,
+      @Nonnull final Precondition precondition,
+      @Nonnull final String field,
+      @Nullable final Object value,
+      final Object... moreFieldsAndValues) {
     verifyNotClosed();
-    BulkCommitBatch bulkCommitBatch = getEligibleBatch();
-    ApiFuture<WriteResult> future =
-        bulkCommitBatch.update(documentReference, precondition, field, value, moreFieldsAndValues);
-    sendReadyBatches();
-    return future;
+    return executeWrite(
+        documentReference,
+        OperationType.UPDATE,
+        new BulkWriterOperationCallback() {
+          public ApiFuture<WriteResult> apply(BulkCommitBatch batch) {
+            return batch.update(documentReference, precondition, field, value, moreFieldsAndValues);
+          }
+        });
   }
 
   /**
@@ -442,22 +597,179 @@ public ApiFuture<WriteResult> update(
    * @param fieldPath The first field to set.
    * @param value The first value to set.
    * @param moreFieldsAndValues String and Object pairs with more fields to be set.
-   * @return An ApiFuture containing the result of the write. Contains an error if the write fails.
+   * @return An ApiFuture containing the result of the write. Contains a {@link BulkWriterException}
+   *     if the write fails.
    */
   @Nonnull
   public ApiFuture<WriteResult> update(
-      @Nonnull DocumentReference documentReference,
-      @Nonnull Precondition precondition,
-      @Nonnull FieldPath fieldPath,
-      @Nullable Object value,
-      Object... moreFieldsAndValues) {
+      @Nonnull final DocumentReference documentReference,
+      @Nonnull final Precondition precondition,
+      @Nonnull final FieldPath fieldPath,
+      @Nullable final Object value,
+      final Object... moreFieldsAndValues) {
     verifyNotClosed();
-    BulkCommitBatch bulkCommitBatch = getEligibleBatch();
-    ApiFuture<WriteResult> future =
-        bulkCommitBatch.update(
-            documentReference, precondition, fieldPath, value, moreFieldsAndValues);
-    sendReadyBatches();
-    return future;
+    return executeWrite(
+        documentReference,
+        OperationType.UPDATE,
+        new BulkWriterOperationCallback() {
+          public ApiFuture<WriteResult> apply(BulkCommitBatch batch) {
+            return batch.update(
+                documentReference, precondition, fieldPath, value, moreFieldsAndValues);
+          }
+        });
+  }
+
+  /**
+   * Schedules the provided write operation and runs the user success callback when the write result
+   * is obtained.
+   */
+  private ApiFuture<WriteResult> executeWrite(
+      final DocumentReference documentReference,
+      final OperationType operationType,
+      final BulkWriterOperationCallback operationCallback) {
+    final SettableApiFuture<Void> operationCompletedFuture = SettableApiFuture.create();
+    writesEnqueued = true;
+
+    ApiFuture<WriteResult> writeResultApiFuture =
+        ApiFutures.transformAsync(
+            enqueueWrite(documentReference, operationType, operationCallback),
+            new ApiAsyncFunction<WriteResult, WriteResult>() {
+              public ApiFuture<WriteResult> apply(final WriteResult result)
+                  throws ExecutionException, InterruptedException {
+
+                invokeUserSuccessCallback(documentReference, result).get();
+                return ApiFutures.immediateFuture(result);
+              }
+            },
+            // This should always run on the `bulkWriterExecutor`. However, we use
+            // `directExecutor()`
+            // here to ensure that the transform code runs immediately after without yielding. Using
+            // `bulkWriterExecutor` instead means that the transform is scheduled after other work
+            // already on the executor, which causes BulkWriter to deadlock.
+            MoreExecutors.directExecutor());
+
+    writeResultApiFuture.addListener(
+        new Runnable() {
+          public void run() {
+            pendingOperations.remove(operationCompletedFuture);
+            operationCompletedFuture.set(null);
+          }
+        },
+        // This should always run on the `bulkWriterExecutor`. However, we use `directExecutor()`
+        // here to ensure that the transform code runs immediately after without yielding. Using
+        // `bulkWriterExecutor` instead means that the transform is scheduled after other work
+        // already on the executor, which causes BulkWriter to deadlock.
+        MoreExecutors.directExecutor());
+
+    pendingOperations.add(operationCompletedFuture);
+    return writeResultApiFuture;
+  }
+
+  /** Used to schedule the write enqueue logic onto the bulkWriterExecutor. */
+  private ApiFuture<WriteResult> enqueueWrite(
+      final DocumentReference documentReference,
+      final OperationType operationType,
+      final BulkWriterOperationCallback operationCallback) {
+    return ApiFutures.transformAsync(
+        ApiFutures.immediateFuture(null),
+        new ApiAsyncFunction<Object, WriteResult>() {
+          public ApiFuture<WriteResult> apply(Object o) throws Exception {
+            return enqueueWriteHelper(documentReference, operationType, operationCallback, 0);
+          }
+        },
+        bulkWriterExecutor);
+  }
+
+  /**
+   * Adds the write to the appropriate batch queue and performs the bulkCommit. This helper method
+   * also catches failures and applies the user error callback and retrying if necessary.
+   */
+  private ApiFuture<WriteResult> enqueueWriteHelper(
+      final DocumentReference documentReference,
+      final OperationType operationType,
+      final BulkWriterOperationCallback operationCallback,
+      final int failedAttempts) {
+    List<BulkCommitBatch> operationBatchQueue = failedAttempts > 0 ? retryBatchQueue : batchQueue;
+    BulkCommitBatch bulkCommitBatch = getEligibleBatch(documentReference, operationBatchQueue);
+
+    // Send ready batches if this is the first attempt. Subsequent retry batches are scheduled
+    // after the initial batch returns.
+    if (failedAttempts == 0) {
+      sendReadyBatches(operationBatchQueue);
+    }
+
+    return ApiFutures.catchingAsync(
+        operationCallback.apply(bulkCommitBatch),
+        FirestoreException.class,
+        new ApiAsyncFunction<FirestoreException, WriteResult>() {
+          public ApiFuture<WriteResult> apply(FirestoreException exception)
+              throws BulkWriterException, ExecutionException, InterruptedException {
+            BulkWriterException bulkWriterException =
+                new BulkWriterException(
+                    exception.getStatus(),
+                    exception.getMessage(),
+                    documentReference,
+                    operationType,
+                    failedAttempts);
+
+            boolean shouldRetry = invokeUserErrorCallback(bulkWriterException).get();
+            logger.log(
+                Level.INFO,
+                String.format(
+                    "Ran error callback on document: %s, error code: %d, shouldRetry: %b",
+                    documentReference.getPath(),
+                    exception.getStatus().getCode().value(),
+                    shouldRetry));
+            if (!shouldRetry) {
+              throw bulkWriterException;
+            } else {
+              return enqueueWriteHelper(
+                  documentReference, operationType, operationCallback, failedAttempts + 1);
+            }
+          }
+        },
+        // This code should always be on the bulkWriterExecutor. However, we use `directExecutor()`
+        // here to ensure that the transform code runs immediately after without yielding. Using
+        // `bulkWriterExecutor` instead means that the transform is scheduled after other work
+        // already on the executor, which causes BulkWriter to deadlock..
+        MoreExecutors.directExecutor());
+  }
+
+  /** Invokes the user error callback on the user callback executor and returns the result. */
+  private SettableApiFuture<Boolean> invokeUserErrorCallback(final BulkWriterException error) {
+    final SettableApiFuture<Boolean> callbackResult = SettableApiFuture.create();
+    errorExecutor.execute(
+        new Runnable() {
+          @Override
+          public void run() {
+            try {
+              boolean shouldRetry = errorListener.onError(error);
+              callbackResult.set(shouldRetry);
+            } catch (Exception e) {
+              callbackResult.setException(e);
+            }
+          }
+        });
+    return callbackResult;
+  }
+
+  /** Invokes the user success callback on the user callback executor. */
+  private ApiFuture<Void> invokeUserSuccessCallback(
+      final DocumentReference documentReference, final WriteResult result) {
+    final SettableApiFuture<Void> callbackResult = SettableApiFuture.create();
+    successExecutor.execute(
+        new Runnable() {
+          @Override
+          public void run() {
+            try {
+              successListener.onResult(documentReference, result);
+              callbackResult.set(null);
+            } catch (Exception e) {
+              callbackResult.setException(e);
+            }
+          }
+        });
+    return callbackResult;
   }
 
   /**
@@ -477,14 +789,50 @@ public ApiFuture<WriteResult> update(
   @Nonnull
   public ApiFuture<Void> flush() {
     verifyNotClosed();
-    final SettableApiFuture<Void> flushComplete = SettableApiFuture.create();
-    List<SettableApiFuture<WriteResult>> writeFutures = new ArrayList<>();
+    final ArrayList<ApiFuture<Void>> pendingOperationsAtFlush =
+        Lists.newArrayList(pendingOperations);
+    return ApiFutures.transformAsync(
+        ApiFutures.immediateFuture(null),
+        new ApiAsyncFunction<Object, Void>() {
+          public ApiFuture<Void> apply(Object o) throws Exception {
+            return performFlush(pendingOperationsAtFlush);
+          }
+        },
+        bulkWriterExecutor);
+  }
+
+  private ApiFuture<Void> performFlush(final List<ApiFuture<Void>> pendingOperations) {
     for (BulkCommitBatch batch : batchQueue) {
       batch.markReadyToSend();
-      writeFutures.addAll(batch.getPendingFutures());
     }
-    sendReadyBatches();
-    ApiFutures.successfulAsList(writeFutures)
+
+    // Send all scheduled operations on the BatchQueue first.
+    sendReadyBatches(batchQueue);
+    final SettableApiFuture<Void> batchQueueFlushComplete = SettableApiFuture.create();
+    final SettableApiFuture<Void> flushComplete = SettableApiFuture.create();
+    ApiFutures.successfulAsList(pendingBatches)
+        .addListener(
+            new Runnable() {
+              public void run() {
+                batchQueueFlushComplete.set(null);
+              }
+            },
+            MoreExecutors.directExecutor());
+
+    // Afterwards, send all accumulated retry operations. Wait until the retryBatchQueue is cleared.
+    // This way, operations scheduled after flush() will not be sent until the retries are done.
+    batchQueueFlushComplete.addListener(
+        new Runnable() {
+          public void run() {
+            for (BulkCommitBatch batch : retryBatchQueue) {
+              batch.markReadyToSend();
+            }
+            sendReadyBatches(retryBatchQueue);
+          }
+        },
+        MoreExecutors.directExecutor());
+
+    ApiFutures.successfulAsList(pendingOperations)
         .addListener(
             new Runnable() {
               public void run() {
@@ -492,13 +840,15 @@ public void run() {
               }
             },
             MoreExecutors.directExecutor());
+
     return flushComplete;
   }
 
   /**
    * Commits all enqueued writes and marks the BulkWriter instance as closed.
    *
-   * <p>After calling `close()`, calling any method wil return an error.
+   * <p>After calling `close()`, calling any method will return an error. Any retries scheduled with
+   * `addWriteErrorListener()` will be run before `close()` completes.
    *
    * <p>This method completes when there are no more pending writes. Calling this method will send
    * all requests.
@@ -515,70 +865,186 @@ private void verifyNotClosed() {
     }
   }
 
+  /**
+   * Attaches a listener that is run every time a BulkWriter operation successfully completes. The
+   * listener will be run before `close()` completes.
+   *
+   * <p>For example, see the sample code: <code>
+   *   BulkWriter bulkWriter = firestore.bulkWriter();
+   *   bulkWriter.addWriteResultListener(
+   *         (DocumentReference documentReference, WriteResult result) -> {
+   *             System.out.println(
+   *                 "Successfully executed write on document: "
+   *                     + documentReference
+   *                     + " at: "
+   *                     + result.getUpdateTime());
+   *           }
+   *         );
+   * </code>
+   *
+   * @param writeResultCallback A callback to be called every time a BulkWriter operation
+   *     successfully completes.
+   */
+  public void addWriteResultListener(WriteResultCallback writeResultCallback) {
+    successListener = writeResultCallback;
+  }
+
+  /**
+   * Attaches a listener that is run every time a BulkWriter operation successfully completes.
+   *
+   * <p>The executor cannot be changed once writes have been enqueued onto the BulkWriter.
+   *
+   * <p>For example, see the sample code: <code>
+   *   BulkWriter bulkWriter = firestore.bulkWriter();
+   *   bulkWriter.addWriteResultListener(
+   *         (DocumentReference documentReference, WriteResult result) -> {
+   *             System.out.println(
+   *                 "Successfully executed write on document: "
+   *                     + documentReference
+   *                     + " at: "
+   *                     + result.getUpdateTime());
+   *           }
+   *         );
+   * </code>
+   *
+   * @param executor The executor to run the provided callback on.
+   * @param writeResultCallback A callback to be called every time a BulkWriter operation
+   *     successfully completes.
+   */
+  public void addWriteResultListener(
+      @Nonnull Executor executor, WriteResultCallback writeResultCallback) {
+    if (writesEnqueued) {
+      throw new IllegalStateException(
+          "The executor cannot be changed once writes have been enqueued.");
+    }
+    successListener = writeResultCallback;
+    successExecutor = executor;
+  }
+
+  /**
+   * Attaches an error handler listener that is run every time a BulkWriter operation fails.
+   *
+   * <p>BulkWriter has a default error handler that retries UNAVAILABLE and ABORTED errors up to a
+   * maximum of 10 failed attempts. When an error handler is specified, the default error handler
+   * will be overwritten.
+   *
+   * <p>For example, see the sample code: <code>
+   *   BulkWriter bulkWriter = firestore.bulkWriter();
+   *   bulkWriter.addWriteErrorListener(
+   *         (BulkWriterException error) -> {
+   *           if (error.getStatus() == Status.UNAVAILABLE
+   *             && error.getFailedAttempts() < MAX_RETRY_ATTEMPTS) {
+   *             return true;
+   *           } else {
+   *             System.out.println("Failed write at document: " + error.getDocumentReference());
+   *             return false;
+   *           }
+   *         }
+   *       );
+   * </code>
+   *
+   * @param onError A callback to be called every time a BulkWriter operation fails. Returning
+   *     `true` will retry the operation. Returning `false` will stop the retry loop.
+   */
+  public void addWriteErrorListener(WriteErrorCallback onError) {
+    errorListener = onError;
+  }
+
+  /**
+   * Attaches an error handler listener that is run every time a BulkWriter operation fails.
+   *
+   * <p>The executor cannot be changed once writes have been enqueued onto the BulkWriter.
+   *
+   * <p>BulkWriter has a default error handler that retries UNAVAILABLE and ABORTED errors up to a
+   * maximum of 10 failed attempts. When an error handler is specified, the default error handler
+   * will be overwritten.
+   *
+   * <p>For example, see the sample code: <code>
+   *   BulkWriter bulkWriter = firestore.bulkWriter();
+   *   bulkWriter.addWriteErrorListener(
+   *         (BulkWriterException error) -> {
+   *           if (error.getStatus() == Status.UNAVAILABLE
+   *             && error.getFailedAttempts() < MAX_RETRY_ATTEMPTS) {
+   *             return true;
+   *           } else {
+   *             System.out.println("Failed write at document: " + error.getDocumentReference());
+   *             return false;
+   *           }
+   *         }
+   *       );
+   * </code>
+   *
+   * @param executor The executor to run the provided callback on.
+   * @param onError A callback to be called every time a BulkWriter operation fails. Returning
+   *     `true` will retry the operation. Returning `false` will stop the retry loop.
+   */
+  public void addWriteErrorListener(@Nonnull Executor executor, WriteErrorCallback onError) {
+    if (writesEnqueued) {
+      throw new IllegalStateException(
+          "The executor cannot be changed once writes have been enqueued.");
+    }
+    errorListener = onError;
+    errorExecutor = executor;
+  }
+
   /**
    * Return the first eligible batch that can hold a write to the provided reference, or creates one
    * if no eligible batches are found.
    */
-  private BulkCommitBatch getEligibleBatch() {
+  private BulkCommitBatch getEligibleBatch(
+      DocumentReference documentReference, List<BulkCommitBatch> batchQueue) {
     if (batchQueue.size() > 0) {
       BulkCommitBatch lastBatch = batchQueue.get(batchQueue.size() - 1);
-      if (lastBatch.getState() == UpdateBuilder.BatchState.OPEN) {
+      if (lastBatch.isOpen() && !lastBatch.has(documentReference)) {
         return lastBatch;
       }
     }
-    return createNewBatch();
+    return createNewBatch(batchQueue);
   }
 
   /**
    * Creates a new batch and adds it to the BatchQueue. If there is already a batch enqueued, sends
    * the batch after a new one is created.
    */
-  private BulkCommitBatch createNewBatch() {
+  private BulkCommitBatch createNewBatch(List<BulkCommitBatch> batchQueue) {
     BulkCommitBatch newBatch = new BulkCommitBatch(firestore, maxBatchSize);
 
     if (batchQueue.size() > 0) {
       batchQueue.get(batchQueue.size() - 1).markReadyToSend();
-      sendReadyBatches();
     }
     batchQueue.add(newBatch);
     return newBatch;
   }
 
   /**
-   * Attempts to send batches starting from the front of the BatchQueue until a batch cannot be
-   * sent.
+   * Attempts to send batches starting from the front of the provided batch queue until a batch
+   * cannot be sent.
    *
    * <p>After a batch is complete, try sending batches again.
    */
-  private void sendReadyBatches() {
-    List<BulkCommitBatch> unsentBatches =
-        FluentIterable.from(batchQueue)
-            .filter(
-                new Predicate<BulkCommitBatch>() {
-                  @Override
-                  public boolean apply(BulkCommitBatch batch) {
-                    return batch.getState() == UpdateBuilder.BatchState.READY_TO_SEND;
-                  }
-                })
-            .toList();
-
+  private void sendReadyBatches(final List<BulkCommitBatch> batchQueue) {
     int index = 0;
-    while (index < unsentBatches.size()
-        && unsentBatches.get(index).state == BatchState.READY_TO_SEND) {
-      final BulkCommitBatch batch = unsentBatches.get(index);
+    while (index < batchQueue.size() && batchQueue.get(index).isReadyToSend()) {
+      final BulkCommitBatch batch = batchQueue.get(index);
+
+      // Future that completes when the current batch or its scheduling attempts completes.
+      final SettableApiFuture<Void> batchCompletedFuture = SettableApiFuture.create();
+      pendingBatches.add(batchCompletedFuture);
 
       // Send the batch if it is under the rate limit, or schedule another attempt after the
       // appropriate timeout.
       long delayMs = rateLimiter.getNextRequestDelayMs(batch.getPendingOperationCount());
       Preconditions.checkState(delayMs != -1, "Batch size should be under capacity");
       if (delayMs == 0) {
-        sendBatch(batch);
+        sendBatch(batch, batchQueue, batchCompletedFuture);
       } else {
-        firestoreExecutor.schedule(
+        bulkWriterExecutor.schedule(
             new Runnable() {
               @Override
               public void run() {
-                sendBatch(batch);
+                sendReadyBatches(batchQueue);
+                pendingBatches.remove(batchCompletedFuture);
+                batchCompletedFuture.set(null);
               }
             },
             delayMs,
@@ -594,11 +1060,10 @@ public void run() {
    * Sends the provided batch and processes the results. After the batch is committed, sends the
    * next group of ready batches.
    */
-  private void sendBatch(final BulkCommitBatch batch) {
-    Preconditions.checkState(
-        batch.state == BatchState.READY_TO_SEND,
-        "The batch should be marked as READY_TO_SEND before committing");
-    batch.state = BatchState.SENT;
+  private void sendBatch(
+      final BulkCommitBatch batch,
+      final List<BulkCommitBatch> batchQueue,
+      final SettableApiFuture<Void> batchCompletedFuture) {
     boolean success = rateLimiter.tryMakeRequest(batch.getPendingOperationCount());
     Preconditions.checkState(success, "Batch should be under rate limit to be sent.");
 
@@ -609,97 +1074,45 @@ public void run() {
             boolean removed = batchQueue.remove(batch);
             Preconditions.checkState(
                 removed, "The batch should be in the BatchQueue." + batchQueue.size());
-            sendReadyBatches();
-          }
-        },
-        MoreExecutors.directExecutor());
-  }
 
-  private ApiFuture<Void> bulkCommit(BulkCommitBatch batch) {
-    return bulkCommit(batch, 0);
-  }
+            if (batchQueue.equals(retryBatchQueue)) {
+              for (BulkCommitBatch batch : retryBatchQueue) {
+                batch.markReadyToSend();
+              }
+            }
 
-  private ApiFuture<Void> bulkCommit(final BulkCommitBatch batch, final int attempt) {
-    final SettableApiFuture<Void> backoffFuture = SettableApiFuture.create();
+            pendingBatches.remove(batchCompletedFuture);
+            batchCompletedFuture.set(null);
 
-    // Add a backoff delay. At first, this is 0.
-    firestoreExecutor.schedule(
-        new Runnable() {
-          @Override
-          public void run() {
-            backoffFuture.set(null);
+            sendReadyBatches(batchQueue);
           }
         },
-        nextAttempt.getRandomizedRetryDelay().toMillis(),
-        TimeUnit.MILLISECONDS);
-
-    return ApiFutures.transformAsync(
-        backoffFuture, new BackoffCallback(batch, attempt), firestoreExecutor);
+        MoreExecutors.directExecutor());
   }
 
-  private class BackoffCallback implements ApiAsyncFunction<Void, Void> {
-    final BulkCommitBatch batch;
-    final int attempt;
-
-    public BackoffCallback(BulkCommitBatch batch, int attempt) {
-      this.batch = batch;
-      this.attempt = attempt;
-    }
-
-    @Override
-    public ApiFuture<Void> apply(Void ignored) {
-
-      return ApiFutures.transformAsync(
-          ApiFutures.catchingAsync(
-              batch.bulkCommit(),
-              Exception.class,
-              new ApiAsyncFunction<Exception, List<BatchWriteResult>>() {
-                public ApiFuture<List<BatchWriteResult>> apply(Exception exception) {
-                  List<BatchWriteResult> results = new ArrayList<>();
-                  // If the BatchWrite RPC fails, map the exception to each individual result.
-                  for (int i = 0; i < batch.getPendingDocumentPaths().size(); ++i) {
-                    results.add(new BatchWriteResult(null, exception));
-                  }
-                  return ApiFutures.immediateFuture(results);
+  private ApiFuture<Void> bulkCommit(final BulkCommitBatch batch) {
+    return ApiFutures.transformAsync(
+        ApiFutures.catchingAsync(
+            batch.bulkCommit(),
+            Exception.class,
+            new ApiAsyncFunction<Exception, List<BatchWriteResult>>() {
+              public ApiFuture<List<BatchWriteResult>> apply(Exception exception) {
+                List<BatchWriteResult> results = new ArrayList<>();
+                // If the BatchWrite RPC fails, map the exception to each individual result.
+                for (int i = 0; i < batch.getPendingOperationCount(); ++i) {
+                  results.add(new BatchWriteResult(null, exception));
                 }
-              },
-              MoreExecutors.directExecutor()),
-          new ProcessBulkCommitCallback(batch, attempt),
-          MoreExecutors.directExecutor());
-    }
-  }
-
-  private class ProcessBulkCommitCallback
-      implements ApiAsyncFunction<List<BatchWriteResult>, Void> {
-    final BulkCommitBatch batch;
-    final int attempt;
-
-    public ProcessBulkCommitCallback(BulkCommitBatch batch, int attempt) {
-      this.batch = batch;
-      this.attempt = attempt;
-    }
-
-    @Override
-    public ApiFuture<Void> apply(List<BatchWriteResult> results) {
-      batch.processResults(results, /* allowRetry= */ true);
-      List<String> remainingOps = batch.getPendingDocumentPaths();
-      if (!remainingOps.isEmpty()) {
-        logger.log(
-            Level.WARNING,
-            String.format(
-                "Current batch failed at retry #%d. Num failures: %d",
-                attempt, remainingOps.size()));
-
-        if (attempt < MAX_RETRY_ATTEMPTS) {
-          nextAttempt = backoff.createNextAttempt(nextAttempt);
-          BulkCommitBatch newBatch = new BulkCommitBatch(firestore, batch);
-          return bulkCommit(newBatch, attempt + 1);
-        } else {
-          batch.processResults(results, /* allowRetry= */ false);
-        }
-      }
-      return ApiFutures.immediateFuture(null);
-    }
+                return ApiFutures.immediateFuture(results);
+              }
+            },
+            MoreExecutors.directExecutor()),
+        new ApiAsyncFunction<List<BatchWriteResult>, Void>() {
+          public ApiFuture<Void> apply(List<BatchWriteResult> results) throws Exception {
+            batch.processResults(results);
+            return ApiFutures.immediateFuture(null);
+          }
+        },
+        MoreExecutors.directExecutor());
   }
 
   @VisibleForTesting
diff --git a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/BulkWriterException.java b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/BulkWriterException.java
new file mode 100644
index 000000000..199ef6671
--- /dev/null
+++ b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/BulkWriterException.java
@@ -0,0 +1,68 @@
+/*
+ * Copyright 2020 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.firestore;
+
+import com.google.cloud.firestore.BulkWriter.OperationType;
+import io.grpc.Status;
+
+/** The error thrown when a BulkWriter operation fails. */
+final class BulkWriterException extends FirestoreException {
+  private final Status status;
+  private final String message;
+  private final DocumentReference documentReference;
+  private final OperationType operationType;
+  private final int failedAttempts;
+
+  public BulkWriterException(
+      Status status,
+      String message,
+      DocumentReference documentReference,
+      OperationType operationType,
+      int failedAttempts) {
+    super(message, status);
+    this.status = status;
+    this.message = message;
+    this.documentReference = documentReference;
+    this.operationType = operationType;
+    this.failedAttempts = failedAttempts;
+  }
+
+  /** @return The status code of the error. */
+  public Status getStatus() {
+    return status;
+  }
+
+  /** @return The error message of the error. */
+  public String getMessage() {
+    return message;
+  }
+
+  /** @return The DocumentReference the operation was performed on. */
+  public DocumentReference getDocumentReference() {
+    return documentReference;
+  }
+
+  /** @return The type of operation performed. */
+  public OperationType getOperationType() {
+    return operationType;
+  }
+
+  /** @return How many times this operation has been attempted unsuccessfully. */
+  public int getFailedAttempts() {
+    return failedAttempts;
+  }
+}
diff --git a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/Firestore.java b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/Firestore.java
index 147ab4b4d..018d92788 100644
--- a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/Firestore.java
+++ b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/Firestore.java
@@ -168,6 +168,9 @@ void getAll(
   @Nonnull
   WriteBatch batch();
 
+  @Nonnull
+  BulkWriter bulkWriter();
+
   /**
    * Returns a FirestoreBundle.Builder {@link FirestoreBundle.Builder} instance using an
    * automatically generated bundle ID. When loaded on clients, client SDKs use the bundle ID and
diff --git a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/FirestoreException.java b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/FirestoreException.java
index 5fabd5e8b..10186bee5 100644
--- a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/FirestoreException.java
+++ b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/FirestoreException.java
@@ -24,10 +24,10 @@
 import javax.annotation.Nullable;
 
 /** A Firestore Service exception. */
-public final class FirestoreException extends BaseGrpcServiceException {
+public class FirestoreException extends BaseGrpcServiceException {
   private Status status;
 
-  private FirestoreException(String reason, Status status) {
+  FirestoreException(String reason, Status status) {
     this(reason, status, null);
   }
 
diff --git a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/FirestoreImpl.java b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/FirestoreImpl.java
index 99d9677c5..aa7b2a510 100644
--- a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/FirestoreImpl.java
+++ b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/FirestoreImpl.java
@@ -39,6 +39,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Random;
+import java.util.concurrent.ScheduledExecutorService;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 
@@ -93,7 +94,7 @@ public WriteBatch batch() {
   }
 
   @Nonnull
-  BulkWriter bulkWriter() {
+  public BulkWriter bulkWriter() {
     return new BulkWriter(this, BulkWriterOptions.builder().setThrottlingEnabled(true).build());
   }
 
@@ -102,6 +103,11 @@ BulkWriter bulkWriter(BulkWriterOptions options) {
     return new BulkWriter(this, options);
   }
 
+  @Nonnull
+  BulkWriter bulkWriter(BulkWriterOptions options, ScheduledExecutorService executor) {
+    return new BulkWriter(this, options, executor);
+  }
+
   @Nonnull
   @Override
   public CollectionReference collection(@Nonnull String collectionPath) {
diff --git a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/Transaction.java b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/Transaction.java
index 26eec365e..a07146ed6 100644
--- a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/Transaction.java
+++ b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/Transaction.java
@@ -80,7 +80,7 @@ public boolean hasTransactionId() {
     return transactionId != null;
   }
 
-  Transaction wrapResult(ApiFuture<WriteResult> result) {
+  Transaction wrapResult(DocumentReference documentReference) {
     return this;
   }
 
diff --git a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/UpdateBuilder.java b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/UpdateBuilder.java
index fc9e04902..217cd1a36 100644
--- a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/UpdateBuilder.java
+++ b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/UpdateBuilder.java
@@ -20,23 +20,14 @@
 import com.google.api.core.ApiFuture;
 import com.google.api.core.ApiFutures;
 import com.google.api.core.InternalExtensionOnly;
-import com.google.api.core.SettableApiFuture;
-import com.google.api.gax.rpc.StatusCode.Code;
-import com.google.cloud.Timestamp;
 import com.google.cloud.firestore.UserDataConverter.EncodingOptions;
-import com.google.cloud.firestore.v1.FirestoreSettings;
-import com.google.common.base.Function;
 import com.google.common.base.Preconditions;
-import com.google.common.collect.FluentIterable;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.util.concurrent.MoreExecutors;
-import com.google.firestore.v1.BatchWriteRequest;
-import com.google.firestore.v1.BatchWriteResponse;
 import com.google.firestore.v1.CommitRequest;
 import com.google.firestore.v1.CommitResponse;
 import com.google.firestore.v1.Write;
 import com.google.protobuf.ByteString;
-import io.grpc.Status;
 import io.opencensus.trace.AttributeValue;
 import io.opencensus.trace.Tracing;
 import java.util.ArrayList;
@@ -44,7 +35,6 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
-import java.util.Set;
 import java.util.SortedSet;
 import java.util.TreeSet;
 import javax.annotation.Nonnull;
@@ -66,62 +56,28 @@ static class WriteOperation {
     }
   }
 
-  /**
-   * Used to represent a pending write operation.
-   *
-   * <p>Contains a pending write's index, document path, and the corresponding result.
-   */
-  static class PendingOperation {
-    int writeBatchIndex;
-    String documentKey;
-    SettableApiFuture<WriteResult> future;
-
-    PendingOperation(
-        int writeBatchIndex, String documentKey, SettableApiFuture<WriteResult> future) {
-      this.writeBatchIndex = writeBatchIndex;
-      this.documentKey = documentKey;
-      this.future = future;
-    }
-  }
-
   final FirestoreImpl firestore;
-  protected final List<WriteOperation> writes;
-  private boolean committed;
-  private final int maxBatchSize;
 
-  protected BatchState state = BatchState.OPEN;
-  protected List<PendingOperation> pendingOperations = new ArrayList<>();
+  private final List<WriteOperation> writes = new ArrayList<>();
 
-  /**
-   * Used to represent the state of batch.
-   *
-   * <p>Writes can only be added while the batch is OPEN. For a batch to be sent, the batch must be
-   * READY_TO_SEND. After a batch is sent, it is marked as SENT.
-   */
-  enum BatchState {
-    OPEN,
-    READY_TO_SEND,
-    SENT,
-  }
+  private boolean committed;
 
-  UpdateBuilder(FirestoreImpl firestore) {
-    this(firestore, BulkWriter.MAX_BATCH_SIZE);
+  boolean isCommitted() {
+    return committed;
   }
 
-  UpdateBuilder(FirestoreImpl firestore, int maxBatchSize) {
+  UpdateBuilder(FirestoreImpl firestore) {
     this.firestore = firestore;
-    this.maxBatchSize = maxBatchSize;
-    this.writes = new ArrayList<>();
   }
 
   /**
    * Wraps the result of the write operation before it is returned.
    *
    * <p>This method is used to generate the return value for all public methods. It allows
-   * operations on Transaction and Writebatch to return the object for chaining, while also allowing
+   * operations on Transaction and WriteBatch to return the object for chaining, while also allowing
    * BulkWriter operations to return the future directly.
    */
-  abstract T wrapResult(ApiFuture<WriteResult> result);
+  abstract T wrapResult(DocumentReference documentReference);
 
   /**
    * Turns a field map that contains field paths into a nested map. Turns {a.b : c} into {a : {b :
@@ -196,12 +152,12 @@ private T performCreate(
 
     writes.add(new WriteOperation(documentReference, write));
 
-    return wrapResult(processOperation(documentReference));
+    return wrapResult(documentReference);
   }
 
   private void verifyNotCommitted() {
     Preconditions.checkState(
-        !committed, "Cannot modify a WriteBatch that has already been committed.");
+        !isCommitted(), "Cannot modify a WriteBatch that has already been committed.");
   }
 
   /**
@@ -326,7 +282,7 @@ private T performSet(
 
     writes.add(new WriteOperation(documentReference, write));
 
-    return wrapResult(processOperation(documentReference));
+    return wrapResult(documentReference);
   }
 
   /** Removes all values in 'fields' that are not specified in 'fieldMask'. */
@@ -595,7 +551,7 @@ public boolean allowTransform() {
     }
     writes.add(new WriteOperation(documentReference, write));
 
-    return wrapResult(processOperation(documentReference));
+    return wrapResult(documentReference);
   }
 
   /**
@@ -633,7 +589,7 @@ private T performDelete(
     }
     writes.add(new WriteOperation(documentReference, write));
 
-    return wrapResult(processOperation(documentReference));
+    return wrapResult(documentReference);
   }
 
   /** Commit the current batch. */
@@ -680,169 +636,17 @@ public List<WriteResult> apply(CommitResponse commitResponse) {
         MoreExecutors.directExecutor());
   }
 
-  /**
-   * Commits all pending operations to the database and verifies all preconditions.
-   *
-   * <p>The writes in the batch are not applied atomically and can be applied out of order.
-   */
-  ApiFuture<List<BatchWriteResult>> bulkCommit() {
-    Tracing.getTracer()
-        .getCurrentSpan()
-        .addAnnotation(
-            TraceUtil.SPAN_NAME_BATCHWRITE,
-            ImmutableMap.of("numDocuments", AttributeValue.longAttributeValue(writes.size())));
-
-    final BatchWriteRequest.Builder request = BatchWriteRequest.newBuilder();
-    request.setDatabase(firestore.getDatabaseName());
-
-    for (WriteOperation writeOperation : writes) {
-      request.addWrites(writeOperation.write);
-    }
-
-    ApiFuture<BatchWriteResponse> response =
-        firestore.sendRequest(request.build(), firestore.getClient().batchWriteCallable());
-
-    return ApiFutures.transform(
-        response,
-        new ApiFunction<BatchWriteResponse, List<BatchWriteResult>>() {
-          @Override
-          public List<BatchWriteResult> apply(BatchWriteResponse batchWriteResponse) {
-            List<com.google.firestore.v1.WriteResult> writeResults =
-                batchWriteResponse.getWriteResultsList();
-
-            List<com.google.rpc.Status> statuses = batchWriteResponse.getStatusList();
-
-            List<BatchWriteResult> result = new ArrayList<>();
-
-            for (int i = 0; i < writeResults.size(); ++i) {
-              com.google.firestore.v1.WriteResult writeResult = writeResults.get(i);
-              com.google.rpc.Status status = statuses.get(i);
-              Status code = Status.fromCodeValue(status.getCode());
-              @Nullable Timestamp updateTime = null;
-              @Nullable Exception exception = null;
-              if (code == Status.OK) {
-                updateTime = Timestamp.fromProto(writeResult.getUpdateTime());
-              } else {
-                exception = FirestoreException.serverRejected(code, status.getMessage());
-              }
-              result.add(new BatchWriteResult(updateTime, exception));
-            }
-
-            return result;
-          }
-        },
-        MoreExecutors.directExecutor());
-  }
-
   /** Checks whether any updates have been queued. */
   boolean isEmpty() {
     return writes.isEmpty();
   }
 
+  List<WriteOperation> getWrites() {
+    return writes;
+  }
+
   /** Get the number of writes. */
   public int getMutationsSize() {
     return writes.size();
   }
-
-  BatchState getState() {
-    return state;
-  }
-
-  List<String> getPendingDocumentPaths() {
-    return FluentIterable.from(pendingOperations)
-        .transform(
-            new Function<PendingOperation, String>() {
-              @Override
-              public String apply(PendingOperation input) {
-                return input.documentKey;
-              }
-            })
-        .toList();
-  }
-
-  List<Integer> getPendingIndexes() {
-    return FluentIterable.from(pendingOperations)
-        .transform(
-            new Function<PendingOperation, Integer>() {
-              @Override
-              public Integer apply(PendingOperation input) {
-                return input.writeBatchIndex;
-              }
-            })
-        .toList();
-  }
-
-  List<SettableApiFuture<WriteResult>> getPendingFutures() {
-    return FluentIterable.from(pendingOperations)
-        .transform(
-            new Function<PendingOperation, SettableApiFuture<WriteResult>>() {
-              @Override
-              public SettableApiFuture<WriteResult> apply(PendingOperation input) {
-                return input.future;
-              }
-            })
-        .toList();
-  }
-
-  int getPendingOperationCount() {
-    return pendingOperations.size();
-  }
-
-  private ApiFuture<WriteResult> processOperation(DocumentReference documentReference) {
-    Preconditions.checkState(state == BatchState.OPEN, "Batch should be OPEN when adding writes");
-    SettableApiFuture<WriteResult> result = SettableApiFuture.create();
-    int nextBatchIndex = getPendingOperationCount();
-    pendingOperations.add(
-        new PendingOperation(nextBatchIndex, documentReference.getPath(), result));
-
-    if (getPendingOperationCount() == maxBatchSize) {
-      state = BatchState.READY_TO_SEND;
-    }
-
-    return result;
-  }
-
-  /**
-   * Resolves the individual operations in the batch with the results and removes the entry from the
-   * pendingOperations map if the result is not retryable.
-   */
-  void processResults(List<BatchWriteResult> results, boolean allowRetry) {
-    List<PendingOperation> newPendingOperations = new ArrayList<>();
-    for (int i = 0; i < results.size(); ++i) {
-      BatchWriteResult result = results.get(i);
-      PendingOperation operation = pendingOperations.get(i);
-
-      if (result.getException() == null) {
-        operation.future.set(new WriteResult(result.getWriteTime()));
-      } else if (!allowRetry || !shouldRetry(result.getException())) {
-        operation.future.setException(result.getException());
-      } else {
-        // Retry the operation if it has not been processed. Store the current index of
-        // pendingOperations to preserve the mapping of this operation's index in the underlying
-        // writes array.
-        newPendingOperations.add(new PendingOperation(i, operation.documentKey, operation.future));
-      }
-    }
-    pendingOperations = newPendingOperations;
-  }
-
-  private boolean shouldRetry(Exception exception) {
-    if (!(exception instanceof FirestoreException)) {
-      return false;
-    }
-    Set<Code> codes = FirestoreSettings.newBuilder().batchWriteSettings().getRetryableCodes();
-    Status status = ((FirestoreException) exception).getStatus();
-    for (Code code : codes) {
-      if (code.equals(Code.valueOf(status.getCode().name()))) {
-        return true;
-      }
-    }
-    return false;
-  }
-
-  void markReadyToSend() {
-    if (state == BatchState.OPEN) {
-      state = BatchState.READY_TO_SEND;
-    }
-  }
 }
diff --git a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/WriteBatch.java b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/WriteBatch.java
index ba69ab420..baedefaa0 100644
--- a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/WriteBatch.java
+++ b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/WriteBatch.java
@@ -41,7 +41,7 @@ public ApiFuture<List<WriteResult>> commit() {
     return super.commit(null);
   }
 
-  WriteBatch wrapResult(ApiFuture<WriteResult> result) {
+  WriteBatch wrapResult(DocumentReference documentReference) {
     return this;
   }
 }
diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/BulkWriterTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/BulkWriterTest.java
index 529f3b4b4..3ab3d5626 100644
--- a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/BulkWriterTest.java
+++ b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/BulkWriterTest.java
@@ -19,8 +19,10 @@
 import static com.google.cloud.firestore.LocalFirestoreHelper.batchWrite;
 import static com.google.cloud.firestore.LocalFirestoreHelper.create;
 import static com.google.cloud.firestore.LocalFirestoreHelper.delete;
+import static com.google.cloud.firestore.LocalFirestoreHelper.map;
 import static com.google.cloud.firestore.LocalFirestoreHelper.set;
 import static com.google.cloud.firestore.LocalFirestoreHelper.update;
+import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
@@ -32,15 +34,21 @@
 import com.google.api.core.SettableApiFuture;
 import com.google.api.gax.rpc.UnaryCallable;
 import com.google.cloud.Timestamp;
+import com.google.cloud.firestore.BulkWriter.WriteErrorCallback;
+import com.google.cloud.firestore.BulkWriter.WriteResultCallback;
 import com.google.cloud.firestore.LocalFirestoreHelper.ResponseStubber;
 import com.google.cloud.firestore.spi.v1.FirestoreRpc;
+import com.google.common.util.concurrent.MoreExecutors;
 import com.google.firestore.v1.BatchWriteRequest;
 import com.google.firestore.v1.BatchWriteResponse;
+import com.google.firestore.v1.Value;
 import com.google.rpc.Code;
 import io.grpc.Status;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.List;
+import java.util.concurrent.Executor;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ScheduledFuture;
@@ -139,7 +147,8 @@ public void hasSetMethod() throws Exception {
         };
     responseStubber.initializeStub(batchWriteCapture, firestoreMock);
 
-    ApiFuture<WriteResult> result = bulkWriter.set(doc1, LocalFirestoreHelper.SINGLE_FIELD_MAP);
+    ApiFuture<WriteResult> result =
+        bulkWriter.set(doc1, (Object) LocalFirestoreHelper.SINGLE_FIELD_MAP);
     bulkWriter.close();
 
     responseStubber.verifyAllRequestsSent();
@@ -152,14 +161,25 @@ public void hasUpdateMethod() throws Exception {
         new ResponseStubber() {
           {
             put(
-                batchWrite(update(LocalFirestoreHelper.SINGLE_FIELD_PROTO, "coll/doc1")),
+                batchWrite(
+                    update(
+                        LocalFirestoreHelper.SINGLE_FIELD_PROTO,
+                        Collections.singletonList("foo"),
+                        com.google.firestore.v1.Precondition.newBuilder()
+                            .setUpdateTime(com.google.protobuf.Timestamp.newBuilder().build())
+                            .build(),
+                        "coll/doc1")),
                 successResponse(2));
           }
         };
     responseStubber.initializeStub(batchWriteCapture, firestoreMock);
 
-    ApiFuture<WriteResult> result = bulkWriter.update(doc1, LocalFirestoreHelper.SINGLE_FIELD_MAP);
-    bulkWriter.close();
+    ApiFuture<WriteResult> result =
+        bulkWriter.update(
+            doc1,
+            LocalFirestoreHelper.SINGLE_FIELD_MAP,
+            Precondition.updatedAt(Timestamp.ofTimeSecondsAndNanos(0, 0)));
+    bulkWriter.flush().get();
 
     responseStubber.verifyAllRequestsSent();
     assertEquals(Timestamp.ofTimeSecondsAndNanos(2, 0), result.get().getUpdateTime());
@@ -175,7 +195,7 @@ public void hasDeleteMethod() throws Exception {
         };
     responseStubber.initializeStub(batchWriteCapture, firestoreMock);
 
-    ApiFuture<WriteResult> result = bulkWriter.delete(doc1);
+    ApiFuture<WriteResult> result = bulkWriter.delete(doc1, Precondition.NONE);
     bulkWriter.close();
 
     responseStubber.verifyAllRequestsSent();
@@ -194,8 +214,9 @@ public void hasCreateMethod() throws Exception {
         };
     responseStubber.initializeStub(batchWriteCapture, firestoreMock);
 
-    ApiFuture<WriteResult> result = bulkWriter.create(doc1, LocalFirestoreHelper.SINGLE_FIELD_MAP);
-    bulkWriter.close();
+    ApiFuture<WriteResult> result =
+        bulkWriter.create(doc1, (Object) LocalFirestoreHelper.SINGLE_FIELD_MAP);
+    bulkWriter.flush().get();
 
     responseStubber.verifyAllRequestsSent();
     assertEquals(Timestamp.ofTimeSecondsAndNanos(2, 0), result.get().getUpdateTime());
@@ -221,8 +242,8 @@ public void surfacesErrors() throws Exception {
       result.get();
       fail("set() should have failed");
     } catch (Exception e) {
-      assertTrue(e.getCause() instanceof FirestoreException);
-      assertEquals(Status.DEADLINE_EXCEEDED, ((FirestoreException) e.getCause()).getStatus());
+      assertTrue(e.getCause() instanceof BulkWriterException);
+      assertEquals(Status.DEADLINE_EXCEEDED, ((BulkWriterException) e.getCause()).getStatus());
     }
   }
 
@@ -240,7 +261,11 @@ public void addsWritesToNewBatchAfterFlush() throws Exception {
                 batchWrite(create(LocalFirestoreHelper.SINGLE_FIELD_PROTO, "coll/doc1")),
                 successResponse(1));
             put(
-                batchWrite(set(LocalFirestoreHelper.SINGLE_FIELD_PROTO, "coll/doc2")),
+                batchWrite(
+                    set(
+                        LocalFirestoreHelper.SINGLE_FIELD_PROTO,
+                        Collections.singletonList("foo"),
+                        "coll/doc2")),
                 successResponse(2));
           }
         };
@@ -248,7 +273,8 @@ public void addsWritesToNewBatchAfterFlush() throws Exception {
 
     ApiFuture<WriteResult> result1 = bulkWriter.create(doc1, LocalFirestoreHelper.SINGLE_FIELD_MAP);
     bulkWriter.flush();
-    ApiFuture<WriteResult> result2 = bulkWriter.set(doc2, LocalFirestoreHelper.SINGLE_FIELD_MAP);
+    ApiFuture<WriteResult> result2 =
+        bulkWriter.set(doc2, LocalFirestoreHelper.SINGLE_FIELD_MAP, SetOptions.merge());
     bulkWriter.close();
 
     responseStubber.verifyAllRequestsSent();
@@ -256,11 +282,6 @@ public void addsWritesToNewBatchAfterFlush() throws Exception {
     assertEquals(Timestamp.ofTimeSecondsAndNanos(2, 0), result2.get().getUpdateTime());
   }
 
-  @Test
-  public void closeResolvesImmediatelyIfNoWrites() throws Exception {
-    bulkWriter.close();
-  }
-
   @Test
   public void cannotCallMethodsAfterClose() throws Exception {
     String expected = "BulkWriter has already been closed.";
@@ -304,15 +325,25 @@ public void cannotCallMethodsAfterClose() throws Exception {
   }
 
   @Test
-  public void canSendWritesToSameDocInSameBatch() throws Exception {
+  public void sendsWritesToSameDocInDifferentBatches() throws Exception {
     ResponseStubber responseStubber =
         new ResponseStubber() {
           {
+            put(
+                batchWrite(set(LocalFirestoreHelper.SINGLE_FIELD_PROTO, "coll/doc1")),
+                successResponse(1));
             put(
                 batchWrite(
-                    set(LocalFirestoreHelper.SINGLE_FIELD_PROTO, "coll/doc1"),
-                    update(LocalFirestoreHelper.SINGLE_FIELD_PROTO, "coll/doc1")),
-                mergeResponses(successResponse(1), successResponse(2)));
+                    update(
+                        map(
+                            "foo",
+                            Value.newBuilder().setStringValue("bar").build(),
+                            "boo",
+                            Value.newBuilder().setStringValue("far").build()),
+                        Arrays.asList("boo", "foo"),
+                        com.google.firestore.v1.Precondition.newBuilder().setExists(true).build(),
+                        "coll/doc1")),
+                successResponse(2));
           }
         };
     responseStubber.initializeStub(batchWriteCapture, firestoreMock);
@@ -320,8 +351,7 @@ public void canSendWritesToSameDocInSameBatch() throws Exception {
     // Create another document reference pointing to the same document.
     DocumentReference sameDoc = firestoreMock.document(doc1.getPath());
     ApiFuture<WriteResult> result1 = bulkWriter.set(doc1, LocalFirestoreHelper.SINGLE_FIELD_MAP);
-    ApiFuture<WriteResult> result2 =
-        bulkWriter.update(sameDoc, LocalFirestoreHelper.SINGLE_FIELD_MAP);
+    ApiFuture<WriteResult> result2 = bulkWriter.update(sameDoc, "foo", "bar", "boo", "far");
     bulkWriter.close();
 
     responseStubber.verifyAllRequestsSent();
@@ -352,6 +382,373 @@ public void sendWritesToDifferentDocsInSameBatch() throws Exception {
     assertEquals(Timestamp.ofTimeSecondsAndNanos(2, 0), result2.get().getUpdateTime());
   }
 
+  @Test
+  public void runsSuccessHandler() throws Exception {
+    ResponseStubber responseStubber =
+        new ResponseStubber() {
+          {
+            put(
+                batchWrite(
+                    create(LocalFirestoreHelper.SINGLE_FIELD_PROTO, "coll/doc1"),
+                    set(LocalFirestoreHelper.SINGLE_FIELD_PROTO, "coll/doc2"),
+                    update(LocalFirestoreHelper.SINGLE_FIELD_PROTO, "coll/doc3"),
+                    delete("coll/doc4")),
+                mergeResponses(
+                    successResponse(1),
+                    successResponse(2),
+                    successResponse(3),
+                    successResponse(4)));
+          }
+        };
+    responseStubber.initializeStub(batchWriteCapture, firestoreMock);
+
+    final List<Integer> writeResults = new ArrayList();
+    DocumentReference doc3 = firestoreMock.document("coll/doc3");
+    DocumentReference doc4 = firestoreMock.document("coll/doc4");
+    bulkWriter.addWriteResultListener(
+        new WriteResultCallback() {
+          public void onResult(DocumentReference documentReference, WriteResult result) {
+            writeResults.add((int) result.getUpdateTime().getSeconds());
+          }
+        });
+    bulkWriter.create(doc1, LocalFirestoreHelper.SINGLE_FIELD_MAP);
+    bulkWriter.set(doc2, LocalFirestoreHelper.SINGLE_FIELD_MAP);
+    bulkWriter.update(doc3, LocalFirestoreHelper.SINGLE_FIELD_MAP);
+    bulkWriter.delete(doc4);
+    bulkWriter.close();
+    assertArrayEquals(new Integer[] {1, 2, 3, 4}, writeResults.toArray());
+  }
+
+  @Test
+  public void retriesFailedOperationsWithGlobalErrorCallback() throws Exception {
+    ResponseStubber responseStubber =
+        new ResponseStubber() {
+          {
+            put(
+                batchWrite(
+                    create(LocalFirestoreHelper.SINGLE_FIELD_PROTO, "coll/doc1"),
+                    set(LocalFirestoreHelper.SINGLE_FIELD_PROTO, "coll/doc2"),
+                    update(LocalFirestoreHelper.SINGLE_FIELD_PROTO, "coll/doc3"),
+                    delete("coll/doc4")),
+                mergeResponses(
+                    successResponse(1),
+                    failedResponse(Code.INTERNAL_VALUE),
+                    failedResponse(Code.INTERNAL_VALUE),
+                    failedResponse(Code.INTERNAL_VALUE)));
+            put(
+                batchWrite(
+                    set(LocalFirestoreHelper.SINGLE_FIELD_PROTO, "coll/doc2"),
+                    update(LocalFirestoreHelper.SINGLE_FIELD_PROTO, "coll/doc3"),
+                    delete("coll/doc4")),
+                mergeResponses(successResponse(2), successResponse(3), successResponse(4)));
+          }
+        };
+    responseStubber.initializeStub(batchWriteCapture, firestoreMock);
+
+    final List<Integer> writeResults = new ArrayList<>();
+    final List<String> operations = new ArrayList<>();
+    DocumentReference doc3 = firestoreMock.document("coll/doc3");
+    DocumentReference doc4 = firestoreMock.document("coll/doc4");
+    Executor userCallbackExecutor = Executors.newSingleThreadExecutor();
+    bulkWriter.addWriteErrorListener(
+        userCallbackExecutor,
+        new WriteErrorCallback() {
+          public boolean onError(BulkWriterException error) {
+            operations.add(error.getOperationType().name());
+            return true;
+          }
+        });
+    bulkWriter.addWriteResultListener(
+        new WriteResultCallback() {
+          public void onResult(DocumentReference documentReference, WriteResult result) {
+            operations.add("SUCCESS");
+            writeResults.add((int) result.getUpdateTime().getSeconds());
+          }
+        });
+    bulkWriter.create(doc1, LocalFirestoreHelper.SINGLE_FIELD_MAP);
+    bulkWriter.set(doc2, LocalFirestoreHelper.SINGLE_FIELD_MAP);
+    bulkWriter.update(doc3, LocalFirestoreHelper.SINGLE_FIELD_MAP);
+    bulkWriter.delete(doc4);
+    bulkWriter.close();
+    assertArrayEquals(new Integer[] {1, 2, 3, 4}, writeResults.toArray());
+    assertArrayEquals(
+        new String[] {"SUCCESS", "SET", "UPDATE", "DELETE", "SUCCESS", "SUCCESS", "SUCCESS"},
+        operations.toArray());
+  }
+
+  @Test
+  public void errorSurfacedEvenWithRetryFunction() throws Exception {
+    ResponseStubber responseStubber =
+        new ResponseStubber() {
+          {
+            put(
+                batchWrite(update(LocalFirestoreHelper.SINGLE_FIELD_PROTO, "coll/doc1")),
+                failedResponse(Code.INTERNAL_VALUE));
+          }
+        };
+    responseStubber.initializeStub(batchWriteCapture, firestoreMock);
+
+    final boolean[] errorListenerCalled = {false};
+    bulkWriter.addWriteErrorListener(
+        new WriteErrorCallback() {
+          public boolean onError(BulkWriterException error) {
+            errorListenerCalled[0] = true;
+            assertEquals(Status.INTERNAL, error.getStatus());
+            return false;
+          }
+        });
+
+    ApiFuture<WriteResult> result = bulkWriter.update(doc1, FieldPath.of("foo"), "bar");
+    bulkWriter.close();
+    assertTrue(errorListenerCalled[0]);
+    try {
+      result.get();
+      fail("Operation should have failed in test");
+    } catch (Exception e) {
+      assertEquals(Status.INTERNAL, ((BulkWriterException) e.getCause()).getStatus());
+    }
+  }
+
+  @Test
+  public void surfacesExceptionsThrownByUserProvidedErrorListener() throws Exception {
+    ResponseStubber responseStubber =
+        new ResponseStubber() {
+          {
+            put(
+                batchWrite(set(LocalFirestoreHelper.SINGLE_FIELD_PROTO, "coll/doc1")),
+                failedResponse(Code.INTERNAL_VALUE));
+          }
+        };
+    responseStubber.initializeStub(batchWriteCapture, firestoreMock);
+
+    bulkWriter.addWriteErrorListener(
+        new WriteErrorCallback() {
+          public boolean onError(BulkWriterException error) {
+            throw new UnsupportedOperationException(
+                "Test code threw UnsupportedOperationException");
+          }
+        });
+
+    ApiFuture<WriteResult> result = bulkWriter.set(doc1, LocalFirestoreHelper.SINGLE_FIELD_MAP);
+    bulkWriter.close();
+    try {
+      result.get();
+      fail("Operation should have failed in test");
+    } catch (Exception e) {
+      assertTrue(e.getMessage().contains("Test code threw UnsupportedOperationException"));
+    }
+  }
+
+  @Test
+  public void writeFailsIfUserProvidedSuccessListenerFails() throws Exception {
+    ResponseStubber responseStubber =
+        new ResponseStubber() {
+          {
+            put(
+                batchWrite(
+                    update(
+                        LocalFirestoreHelper.SINGLE_FIELD_PROTO,
+                        Collections.singletonList("foo"),
+                        com.google.firestore.v1.Precondition.newBuilder()
+                            .setUpdateTime(com.google.protobuf.Timestamp.newBuilder().build())
+                            .build(),
+                        "coll/doc1")),
+                successResponse(1));
+          }
+        };
+    responseStubber.initializeStub(batchWriteCapture, firestoreMock);
+
+    bulkWriter.addWriteResultListener(
+        new WriteResultCallback() {
+          public void onResult(DocumentReference documentReference, WriteResult result) {
+            throw new UnsupportedOperationException(
+                "Test code threw UnsupportedOperationException");
+          }
+        });
+
+    ApiFuture<WriteResult> result =
+        bulkWriter.update(
+            doc1, Precondition.updatedAt(Timestamp.ofTimeSecondsAndNanos(0, 0)), "foo", "bar");
+    bulkWriter.close();
+    try {
+      result.get();
+      fail("Operation should have failed in test");
+    } catch (Exception e) {
+      assertTrue(e.getMessage().contains("Test code threw UnsupportedOperationException"));
+    }
+  }
+
+  @Test
+  public void cannotChangeExecutorOnceWriteEnqueued() throws Exception {
+    bulkWriter.set(doc1, LocalFirestoreHelper.SINGLE_FIELD_MAP);
+    try {
+      bulkWriter.addWriteResultListener(
+          MoreExecutors.directExecutor(),
+          new WriteResultCallback() {
+            public void onResult(DocumentReference documentReference, WriteResult result) {}
+          });
+      fail("Operation should have failed in test");
+    } catch (Exception e) {
+      assertTrue(
+          e.getMessage().contains("The executor cannot be changed once writes have been enqueued"));
+    }
+
+    try {
+      bulkWriter.addWriteErrorListener(
+          MoreExecutors.directExecutor(),
+          new WriteErrorCallback() {
+            public boolean onError(BulkWriterException error) {
+              return false;
+            }
+          });
+      fail("Operation should have failed in test");
+    } catch (Exception e) {
+      assertTrue(
+          e.getMessage().contains("The executor cannot be changed once writes have been enqueued"));
+    }
+  }
+
+  @Test
+  public void retriesMultipleTimes() throws Exception {
+    ResponseStubber responseStubber =
+        new ResponseStubber() {
+          {
+            put(
+                batchWrite(set(LocalFirestoreHelper.SINGLE_FIELD_PROTO, "coll/doc1")),
+                failedResponse(Code.INTERNAL_VALUE));
+            put(
+                batchWrite(set(LocalFirestoreHelper.SINGLE_FIELD_PROTO, "coll/doc1")),
+                failedResponse(Code.INTERNAL_VALUE));
+            put(
+                batchWrite(set(LocalFirestoreHelper.SINGLE_FIELD_PROTO, "coll/doc1")),
+                failedResponse(Code.INTERNAL_VALUE));
+            put(
+                batchWrite(set(LocalFirestoreHelper.SINGLE_FIELD_PROTO, "coll/doc1")),
+                successResponse(1));
+          }
+        };
+    responseStubber.initializeStub(batchWriteCapture, firestoreMock);
+
+    bulkWriter.addWriteErrorListener(
+        new WriteErrorCallback() {
+          public boolean onError(BulkWriterException error) {
+            return true;
+          }
+        });
+
+    ApiFuture<WriteResult> result1 = bulkWriter.set(doc1, LocalFirestoreHelper.SINGLE_FIELD_MAP);
+    bulkWriter.close();
+    assertEquals(Timestamp.ofTimeSecondsAndNanos(1, 0), result1.get().getUpdateTime());
+  }
+
+  @Test
+  public void retriesMaintainCorrectWriteResolutionOrdering() throws Exception {
+    ResponseStubber responseStubber =
+        new ResponseStubber() {
+          {
+            put(
+                batchWrite(set(LocalFirestoreHelper.SINGLE_FIELD_PROTO, "coll/doc1")),
+                failedResponse(Code.INTERNAL_VALUE));
+            put(
+                batchWrite(set(LocalFirestoreHelper.SINGLE_FIELD_PROTO, "coll/doc1")),
+                successResponse(1));
+            put(
+                batchWrite(set(LocalFirestoreHelper.SINGLE_FIELD_PROTO, "coll/doc2")),
+                successResponse(2));
+          }
+        };
+    responseStubber.initializeStub(batchWriteCapture, firestoreMock);
+    ScheduledExecutorService executor = Executors.newSingleThreadScheduledExecutor();
+
+    // Use separate futures to track listener completion since the callbacks are run on a different
+    // thread than the BulkWriter operations.
+    final SettableApiFuture<Void> flushComplete = SettableApiFuture.create();
+    final SettableApiFuture<Void> closeComplete = SettableApiFuture.create();
+
+    final List<String> operations = new ArrayList<>();
+    bulkWriter.addWriteErrorListener(
+        new WriteErrorCallback() {
+          public boolean onError(BulkWriterException error) {
+            return true;
+          }
+        });
+
+    bulkWriter
+        .set(doc1, LocalFirestoreHelper.SINGLE_FIELD_MAP)
+        .addListener(
+            new Runnable() {
+              public void run() {
+                operations.add("BEFORE_FLUSH");
+              }
+            },
+            executor);
+    bulkWriter
+        .flush()
+        .addListener(
+            new Runnable() {
+              public void run() {
+                operations.add("FLUSH");
+                flushComplete.set(null);
+              }
+            },
+            executor);
+    bulkWriter
+        .set(doc2, LocalFirestoreHelper.SINGLE_FIELD_MAP)
+        .addListener(
+            new Runnable() {
+              public void run() {
+                operations.add("AFTER_FLUSH");
+                closeComplete.set(null);
+              }
+            },
+            executor);
+    flushComplete.get();
+
+    // Verify that the 2nd operation did not complete as a result of the flush() call.
+    assertArrayEquals(new String[] {"BEFORE_FLUSH", "FLUSH"}, operations.toArray());
+    bulkWriter.close();
+    closeComplete.get();
+    assertArrayEquals(new String[] {"BEFORE_FLUSH", "FLUSH", "AFTER_FLUSH"}, operations.toArray());
+  }
+
+  @Test
+  public void returnsTheErrorIfNoRetrySpecified() throws Exception {
+    ResponseStubber responseStubber =
+        new ResponseStubber() {
+          {
+            put(
+                batchWrite(set(LocalFirestoreHelper.SINGLE_FIELD_PROTO, "coll/doc1")),
+                failedResponse(Code.INTERNAL_VALUE));
+            put(
+                batchWrite(set(LocalFirestoreHelper.SINGLE_FIELD_PROTO, "coll/doc1")),
+                failedResponse(Code.INTERNAL_VALUE));
+            put(
+                batchWrite(set(LocalFirestoreHelper.SINGLE_FIELD_PROTO, "coll/doc1")),
+                failedResponse(Code.INTERNAL_VALUE));
+            put(
+                batchWrite(set(LocalFirestoreHelper.SINGLE_FIELD_PROTO, "coll/doc1")),
+                failedResponse(Code.INTERNAL_VALUE));
+          }
+        };
+    responseStubber.initializeStub(batchWriteCapture, firestoreMock);
+
+    bulkWriter.addWriteErrorListener(
+        new WriteErrorCallback() {
+          public boolean onError(BulkWriterException error) {
+            return error.getFailedAttempts() < 3;
+          }
+        });
+
+    ApiFuture<WriteResult> result1 = bulkWriter.set(doc1, LocalFirestoreHelper.SINGLE_FIELD_MAP);
+    bulkWriter.close();
+    try {
+      result1.get();
+      fail("Operation should have failed");
+    } catch (Exception e) {
+      assertEquals(Status.INTERNAL, ((BulkWriterException) e.getCause()).getStatus());
+    }
+  }
+
   @Test
   public void sendBatchesWhenSizeLimitIsReached() throws Exception {
     ResponseStubber responseStubber =
@@ -360,7 +757,13 @@ public void sendBatchesWhenSizeLimitIsReached() throws Exception {
             put(
                 batchWrite(
                     set(LocalFirestoreHelper.SINGLE_FIELD_PROTO, "coll/doc1"),
-                    update(LocalFirestoreHelper.SINGLE_FIELD_PROTO, "coll/doc2"),
+                    update(
+                        LocalFirestoreHelper.SINGLE_FIELD_PROTO,
+                        Collections.singletonList("foo"),
+                        com.google.firestore.v1.Precondition.newBuilder()
+                            .setUpdateTime(com.google.protobuf.Timestamp.newBuilder().build())
+                            .build(),
+                        "coll/doc2"),
                     create(LocalFirestoreHelper.SINGLE_FIELD_PROTO, "coll/doc3")),
                 mergeResponses(successResponse(1), successResponse(2), successResponse(3)));
           }
@@ -369,7 +772,12 @@ public void sendBatchesWhenSizeLimitIsReached() throws Exception {
 
     bulkWriter.setMaxBatchSize(3);
     ApiFuture<WriteResult> result1 = bulkWriter.set(doc1, LocalFirestoreHelper.SINGLE_FIELD_MAP);
-    ApiFuture<WriteResult> result2 = bulkWriter.update(doc2, LocalFirestoreHelper.SINGLE_FIELD_MAP);
+    ApiFuture<WriteResult> result2 =
+        bulkWriter.update(
+            doc2,
+            Precondition.updatedAt(Timestamp.ofTimeSecondsAndNanos(0, 0)),
+            FieldPath.of("foo"),
+            "bar");
     ApiFuture<WriteResult> result3 =
         bulkWriter.create(
             firestoreMock.document("coll/doc3"), LocalFirestoreHelper.SINGLE_FIELD_MAP);
@@ -392,7 +800,7 @@ public void retriesIndividualWritesThatFailWithAbortedOrUnavailable() throws Exc
             put(
                 batchWrite(
                     set(LocalFirestoreHelper.SINGLE_FIELD_PROTO, "coll/doc1"),
-                    set(LocalFirestoreHelper.UPDATED_FIELD_PROTO, "coll/doc1"),
+                    set(LocalFirestoreHelper.UPDATED_FIELD_PROTO, "coll/doc2"),
                     set(LocalFirestoreHelper.SINGLE_FIELD_PROTO, "coll/doc3")),
                 mergeResponses(
                     failedResponse(),
@@ -400,7 +808,7 @@ public void retriesIndividualWritesThatFailWithAbortedOrUnavailable() throws Exc
                     failedResponse(Code.ABORTED_VALUE)));
             put(
                 batchWrite(
-                    set(LocalFirestoreHelper.UPDATED_FIELD_PROTO, "coll/doc1"),
+                    set(LocalFirestoreHelper.UPDATED_FIELD_PROTO, "coll/doc2"),
                     set(LocalFirestoreHelper.SINGLE_FIELD_PROTO, "coll/doc3")),
                 mergeResponses(successResponse(2), failedResponse(Code.ABORTED_VALUE)));
             put(
@@ -413,7 +821,7 @@ public void retriesIndividualWritesThatFailWithAbortedOrUnavailable() throws Exc
     // Test writes to the same document in order to verify that retry logic unaffected by document
     // key.
     ApiFuture<WriteResult> result1 = bulkWriter.set(doc1, LocalFirestoreHelper.SINGLE_FIELD_MAP);
-    ApiFuture<WriteResult> result2 = bulkWriter.set(doc1, LocalFirestoreHelper.UPDATED_FIELD_MAP);
+    ApiFuture<WriteResult> result2 = bulkWriter.set(doc2, LocalFirestoreHelper.UPDATED_FIELD_MAP);
     ApiFuture<WriteResult> result3 =
         bulkWriter.set(firestoreMock.document("coll/doc3"), LocalFirestoreHelper.SINGLE_FIELD_MAP);
     bulkWriter.close();
@@ -422,8 +830,8 @@ public void retriesIndividualWritesThatFailWithAbortedOrUnavailable() throws Exc
       result1.get();
       fail("set() should have failed");
     } catch (Exception e) {
-      assertTrue(e.getCause() instanceof FirestoreException);
-      assertEquals(Status.DEADLINE_EXCEEDED, ((FirestoreException) e.getCause()).getStatus());
+      assertTrue(e.getCause() instanceof BulkWriterException);
+      assertEquals(Status.DEADLINE_EXCEEDED, ((BulkWriterException) e.getCause()).getStatus());
     }
     assertEquals(Timestamp.ofTimeSecondsAndNanos(2, 0), result2.get().getUpdateTime());
     assertEquals(Timestamp.ofTimeSecondsAndNanos(3, 0), result3.get().getUpdateTime());
@@ -438,7 +846,10 @@ public void writesCompleteInCorrectOrderBeforeFlush() throws Exception {
           {
             put(
                 batchWrite(
-                    set(LocalFirestoreHelper.SINGLE_FIELD_PROTO, "coll/doc1"),
+                    set(
+                        LocalFirestoreHelper.SINGLE_FIELD_PROTO,
+                        Collections.singletonList("foo"),
+                        "coll/doc1"),
                     set(LocalFirestoreHelper.SINGLE_FIELD_PROTO, "coll/doc2")),
                 mergeResponses(successResponse(1), failedResponse(Code.ABORTED_VALUE)));
             put(
@@ -452,7 +863,7 @@ public void writesCompleteInCorrectOrderBeforeFlush() throws Exception {
     final SettableApiFuture<Void> flushComplete = SettableApiFuture.create();
 
     bulkWriter
-        .set(doc1, LocalFirestoreHelper.SINGLE_FIELD_MAP)
+        .set(doc1, (Object) LocalFirestoreHelper.SINGLE_FIELD_MAP, SetOptions.merge())
         .addListener(
             new Runnable() {
               public void run() {
@@ -513,29 +924,42 @@ public void flushCompletesWhenAllWritesComplete() throws Exception {
   }
 
   @Test
-  public void doesNotSendBatchesIfDoingSoExceedsRateLimit() {
-    final boolean[] timeoutCalled = {false};
+  public void doesNotSendBatchesIfDoingSoExceedsRateLimit() throws Exception {
+    // This future is completed when the BulkWriter schedules a timeout. This test waits on the
+    // future at the end of the test to ensure that the timeout was called.
+    final SettableApiFuture<Void> timeoutCalledFuture = SettableApiFuture.create();
+
     final ScheduledExecutorService timeoutExecutor =
         new ScheduledThreadPoolExecutor(1) {
           @Override
           @Nonnull
           public ScheduledFuture<?> schedule(Runnable command, long delay, TimeUnit unit) {
             if (delay > 0) {
-              timeoutCalled[0] = true;
+              timeoutCalledFuture.set(null);
             }
             return super.schedule(command, 0, TimeUnit.MILLISECONDS);
           }
         };
-    doReturn(timeoutExecutor).when(firestoreRpc).getExecutor();
+
+    // Stub responses for the BulkWriter batches.
+    ResponseStubber responseStubber =
+        new ResponseStubber() {
+          {
+            put(
+                batchWrite(set(LocalFirestoreHelper.SINGLE_FIELD_PROTO, "coll/doc")),
+                successResponse(5));
+          }
+        };
+    responseStubber.initializeStub(batchWriteCapture, firestoreMock);
     BulkWriter bulkWriter =
-        firestoreMock.bulkWriter(BulkWriterOptions.builder().setInitialOpsPerSecond(5).build());
+        firestoreMock.bulkWriter(
+            BulkWriterOptions.builder().setInitialOpsPerSecond(5).build(), timeoutExecutor);
 
     for (int i = 0; i < 600; ++i) {
-      bulkWriter.set(firestoreMock.document("coll/doc" + i), LocalFirestoreHelper.SINGLE_FIELD_MAP);
+      bulkWriter.set(firestoreMock.document("coll/doc"), LocalFirestoreHelper.SINGLE_FIELD_MAP);
     }
     bulkWriter.flush();
-
-    assertTrue(timeoutCalled[0]);
+    timeoutCalledFuture.get();
   }
 
   @Test
@@ -659,7 +1083,7 @@ public ApiFuture<Object> answer(InvocationOnMock mock) {
             batchWriteCapture.capture(),
             Matchers.<UnaryCallable<BatchWriteRequest, BatchWriteResponse>>any());
     ApiFuture<WriteResult> result = bulkWriter.set(doc1, LocalFirestoreHelper.SINGLE_FIELD_MAP);
-    bulkWriter.close();
+    bulkWriter.flush().get();
 
     try {
       result.get();
diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/ITBulkWriterTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/ITBulkWriterTest.java
new file mode 100644
index 000000000..0fb7ba585
--- /dev/null
+++ b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/ITBulkWriterTest.java
@@ -0,0 +1,216 @@
+/*
+ * Copyright 2020 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.firestore;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
+
+import com.google.api.core.ApiFuture;
+import com.google.cloud.Timestamp;
+import com.google.cloud.firestore.BulkWriter.WriteErrorCallback;
+import com.google.cloud.firestore.BulkWriter.WriteResultCallback;
+import com.google.common.base.Preconditions;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.Executor;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ThreadFactory;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TestName;
+
+public class ITBulkWriterTest {
+  private Firestore firestore;
+  private CollectionReference randomColl;
+  private DocumentReference randomDoc;
+
+  @Rule public TestName testName = new TestName();
+
+  @Before
+  public void before() {
+    FirestoreOptions firestoreOptions = FirestoreOptions.newBuilder().build();
+    firestore = firestoreOptions.getService();
+    randomColl =
+        firestore.collection(
+            String.format("java-%s-%s", testName.getMethodName(), LocalFirestoreHelper.autoId()));
+    randomDoc = randomColl.document();
+  }
+
+  @After
+  public void after() throws Exception {
+    Preconditions.checkNotNull(
+        firestore,
+        "Error instantiating Firestore. Check that the service account credentials were properly set.");
+    firestore.close();
+  }
+
+  @Test
+  public void bulkWriterCreate() throws Exception {
+    DocumentReference docRef = randomColl.document();
+    firestore.bulkWriter();
+
+    BulkWriter writer = firestore.bulkWriter();
+    ApiFuture<WriteResult> result =
+        writer.create(docRef, Collections.singletonMap("foo", (Object) "bar"));
+    writer.close();
+
+    assertNotNull(result.get().getUpdateTime());
+    DocumentSnapshot snapshot = docRef.get().get();
+    assertEquals("bar", snapshot.get("foo"));
+  }
+
+  @Test
+  public void bulkWriterCreateAddsPrecondition() throws Exception {
+    DocumentReference docRef = randomColl.document();
+    docRef.set(Collections.singletonMap("foo", (Object) "bar")).get();
+
+    BulkWriter writer = firestore.bulkWriter();
+    ApiFuture<WriteResult> result =
+        writer.create(docRef, Collections.singletonMap("foo", (Object) "bar"));
+    writer.close();
+
+    try {
+      result.get();
+      fail("Create operation should have thrown exception");
+    } catch (Exception e) {
+      assertTrue(e.getMessage().contains("Document already exists"));
+    }
+  }
+
+  @Test
+  public void bulkWriterSet() throws Exception {
+    DocumentReference docRef = randomColl.document();
+
+    BulkWriter writer = firestore.bulkWriter();
+    ApiFuture<WriteResult> result =
+        writer.set(docRef, Collections.singletonMap("foo", (Object) "bar"));
+    writer.close();
+
+    assertNotNull(result.get().getUpdateTime());
+    DocumentSnapshot snapshot = docRef.get().get();
+    assertEquals("bar", snapshot.get("foo"));
+  }
+
+  @Test
+  public void bulkWriterUpdate() throws Exception {
+    DocumentReference docRef = randomColl.document();
+    docRef.set(Collections.singletonMap("foo", "oldValue")).get();
+
+    BulkWriter writer = firestore.bulkWriter();
+    ApiFuture<WriteResult> result = writer.update(docRef, "foo", "newValue");
+    writer.close();
+
+    assertNotNull(result.get().getUpdateTime());
+    DocumentSnapshot snapshot = docRef.get().get();
+    assertEquals("newValue", snapshot.get("foo"));
+  }
+
+  @Test
+  public void bulkWriterUpdateAddsPrecondition() throws Exception {
+    DocumentReference docRef = randomColl.document();
+
+    BulkWriter writer = firestore.bulkWriter();
+    ApiFuture<WriteResult> result = writer.update(docRef, "foo", "newValue");
+    writer.close();
+
+    try {
+      result.get();
+      fail("Update operation should have thrown exception");
+    } catch (Exception e) {
+      assertTrue(e.getMessage().matches(".* No (document|entity) to update.*"));
+    }
+  }
+
+  @Test
+  public void bulkWriterDelete() throws Exception {
+    DocumentReference docRef = randomColl.document();
+    docRef.set(Collections.singletonMap("foo", "oldValue")).get();
+
+    BulkWriter writer = firestore.bulkWriter();
+    ApiFuture<WriteResult> result = writer.delete(docRef);
+    writer.close();
+
+    assertNotNull(result.get().getUpdateTime());
+    // TODO(b/158502664): Remove this check once we can get write times.
+    assertEquals(Timestamp.ofTimeSecondsAndNanos(0, 0), result.get().getUpdateTime());
+    DocumentSnapshot snapshot = docRef.get().get();
+    assertNull(snapshot.get("foo"));
+  }
+
+  @Test
+  public void bulkWriterOnResult() throws Exception {
+    class NamedThreadFactory implements ThreadFactory {
+      public Thread newThread(Runnable r) {
+        return new Thread(r, "bulkWriterSuccess");
+      }
+    }
+    Executor executor = Executors.newSingleThreadExecutor(new NamedThreadFactory());
+    final List<String> operations = new ArrayList<>();
+
+    BulkWriter writer = firestore.bulkWriter();
+    writer.addWriteResultListener(
+        executor,
+        new WriteResultCallback() {
+          public void onResult(DocumentReference documentReference, WriteResult result) {
+            operations.add("operation");
+            assertTrue(Thread.currentThread().getName().contains("bulkWriterSuccess"));
+          }
+        });
+    writer.set(randomDoc, Collections.singletonMap("foo", "bar"));
+    writer.flush().get();
+    assertEquals("operation", operations.get(0));
+  }
+
+  @Test
+  public void bulkWriterOnError() throws Exception {
+    class NamedThreadFactory implements ThreadFactory {
+      public Thread newThread(Runnable r) {
+        return new Thread(r, "bulkWriterException");
+      }
+    }
+    Executor executor = Executors.newSingleThreadExecutor(new NamedThreadFactory());
+    final List<String> operations = new ArrayList<>();
+
+    BulkWriter writer = firestore.bulkWriter();
+    writer.addWriteErrorListener(
+        executor,
+        new WriteErrorCallback() {
+          public boolean onError(BulkWriterException error) {
+            operations.add("operation-error");
+            assertTrue(Thread.currentThread().getName().contains("bulkWriterException"));
+            return false;
+          }
+        });
+
+    writer.addWriteResultListener(
+        executor,
+        new WriteResultCallback() {
+          public void onResult(DocumentReference documentReference, WriteResult result) {
+            fail("The success listener shouldn't be called");
+          }
+        });
+    writer.update(randomDoc, "foo", "bar");
+    writer.flush().get();
+    assertEquals("operation-error", operations.get(0));
+  }
+}
diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/WriteBatchTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/WriteBatchTest.java
index ecf2930c3..7dace6829 100644
--- a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/WriteBatchTest.java
+++ b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/WriteBatchTest.java
@@ -18,7 +18,6 @@
 
 import static com.google.cloud.firestore.LocalFirestoreHelper.UPDATED_SINGLE_FIELD_PROTO;
 import static com.google.cloud.firestore.LocalFirestoreHelper.UPDATE_SINGLE_FIELD_OBJECT;
-import static com.google.cloud.firestore.LocalFirestoreHelper.batchWrite;
 import static com.google.cloud.firestore.LocalFirestoreHelper.commit;
 import static com.google.cloud.firestore.LocalFirestoreHelper.commitResponse;
 import static com.google.cloud.firestore.LocalFirestoreHelper.create;
@@ -27,19 +26,15 @@
 import static com.google.cloud.firestore.LocalFirestoreHelper.set;
 import static com.google.cloud.firestore.LocalFirestoreHelper.update;
 import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertNull;
 import static org.mockito.Mockito.doReturn;
 
-import com.google.api.core.ApiFutures;
 import com.google.api.gax.rpc.UnaryCallable;
 import com.google.cloud.Timestamp;
 import com.google.cloud.firestore.spi.v1.FirestoreRpc;
 import com.google.firestore.v1.BatchWriteRequest;
-import com.google.firestore.v1.BatchWriteResponse;
 import com.google.firestore.v1.CommitRequest;
 import com.google.firestore.v1.CommitResponse;
 import com.google.firestore.v1.Write;
-import io.grpc.Status;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
@@ -336,40 +331,4 @@ public void deleteDocument() throws Exception {
     CommitRequest commitRequest = commitCapture.getValue();
     assertEquals(commit(writes.toArray(new Write[] {})), commitRequest);
   }
-
-  @Test
-  public void bulkCommit() throws Exception {
-    BatchWriteResponse.Builder response = BatchWriteResponse.newBuilder();
-    response.addWriteResultsBuilder().getUpdateTimeBuilder().setNanos(1);
-    response.addWriteResultsBuilder();
-    response.addStatusBuilder().build();
-    response.addStatusBuilder().setCode(14).build();
-    doReturn(ApiFutures.immediateFuture(response.build()))
-        .when(firestoreMock)
-        .sendRequest(
-            batchWriteCapture.capture(),
-            Matchers.<UnaryCallable<BatchWriteRequest, BatchWriteResponse>>any());
-
-    List<Write> writes = new ArrayList<>();
-    batch.set(documentReference, LocalFirestoreHelper.SINGLE_FIELD_MAP);
-    writes.add(set(LocalFirestoreHelper.SINGLE_FIELD_PROTO));
-
-    batch.create(documentReference, LocalFirestoreHelper.SINGLE_FIELD_MAP);
-    writes.add(create(LocalFirestoreHelper.SINGLE_FIELD_PROTO));
-
-    assertEquals(2, batch.getMutationsSize());
-
-    batch.markReadyToSend();
-    List<BatchWriteResult> batchWriteResults = batch.bulkCommit().get();
-
-    assertNull(batchWriteResults.get(0).getException());
-    assertEquals(Timestamp.ofTimeSecondsAndNanos(0, 1), batchWriteResults.get(0).getWriteTime());
-    assertEquals(
-        Status.UNAVAILABLE,
-        ((FirestoreException) batchWriteResults.get(1).getException()).getStatus());
-    assertNull(batchWriteResults.get(1).getWriteTime());
-
-    BatchWriteRequest batchWriteRequest = batchWriteCapture.getValue();
-    assertEquals(batchWrite(writes.toArray(new Write[] {})), batchWriteRequest);
-  }
 }