Adding .close() to Firestore Service (googleapis#2807)

schmidt-sebastian · garrettjonesgoogle · commit 80e304f8199c · 2018-01-25T16:42:55.000-08:00
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
@@ -22,7 +22,7 @@
 import javax.annotation.Nonnull;
 
 /** Represents a Firestore Database and is the entry point for all Firestore operations */
-public interface Firestore extends Service<FirestoreOptions> {
+public interface Firestore extends Service<FirestoreOptions>, AutoCloseable {
 
   /**
    * Gets a {@link CollectionReference} that refers to the collection at the specified path.
@@ -91,4 +91,12 @@ <T> ApiFuture<T> runTransaction(
    */
   @Nonnull
   WriteBatch batch();
+
+  /**
+   * Closes the gRPC channels associated with this instance and frees up their resources. This
+   * method blocks until all channels are closed. Once this method is called, this Firestore client
+   * is no longer usable.
+   */
+  @Override
+  void close() throws Exception;
 }
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
@@ -52,20 +52,12 @@ class FirestoreImpl implements Firestore {
   private static final String AUTO_ID_ALPHABET =
       "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789";
 
-  /** Creates a pseudo-random 20-character ID that can be used for Firestore documents. */
-  static String autoId() {
-    StringBuilder builder = new StringBuilder();
-    int maxRandom = AUTO_ID_ALPHABET.length();
-    for (int i = 0; i < AUTO_ID_LENGTH; i++) {
-      builder.append(AUTO_ID_ALPHABET.charAt(RANDOM.nextInt(maxRandom)));
-    }
-    return builder.toString();
-  }
-
   private final FirestoreRpc firestoreClient;
   private final FirestoreOptions firestoreOptions;
   private final ResourcePath databasePath;
 
+  private boolean closed;
+
   FirestoreImpl(FirestoreOptions options) {
     this(options, options.getFirestoreRpc());
   }
@@ -81,6 +73,16 @@ static String autoId() {
         ResourcePath.create(DatabaseRootName.of(options.getProjectId(), options.getDatabaseId()));
   }
 
+  /** Creates a pseudo-random 20-character ID that can be used for Firestore documents. */
+  static String autoId() {
+    StringBuilder builder = new StringBuilder();
+    int maxRandom = AUTO_ID_ALPHABET.length();
+    for (int i = 0; i < AUTO_ID_LENGTH; i++) {
+      builder.append(AUTO_ID_ALPHABET.charAt(RANDOM.nextInt(maxRandom)));
+    }
+    return builder.toString();
+  }
+
   @Nonnull
   @Override
   public WriteBatch batch() {
@@ -324,6 +326,7 @@ FirestoreRpc getClient() {
   /** Request funnel for all read/write requests. */
   <RequestT, ResponseT> ApiFuture<ResponseT> sendRequest(
       RequestT requestT, UnaryCallable<RequestT, ResponseT> callable) {
+    Preconditions.checkState(!closed, "Firestore client has already been closed");
     return callable.futureCall(requestT);
   }
 
@@ -332,18 +335,26 @@ <RequestT, ResponseT> void streamRequest(
       RequestT requestT,
       ApiStreamObserver<ResponseT> responseObserverT,
       ServerStreamingCallable<RequestT, ResponseT> callable) {
+    Preconditions.checkState(!closed, "Firestore client has already been closed");
     callable.serverStreamingCall(requestT, responseObserverT);
   }
 
   /** Request funnel for all bidirectional streaming requests. */
   <RequestT, ResponseT> ApiStreamObserver<RequestT> streamRequest(
       ApiStreamObserver<ResponseT> responseObserverT,
       BidiStreamingCallable<RequestT, ResponseT> callable) {
+    Preconditions.checkState(!closed, "Firestore client has already been closed");
     return callable.bidiStreamingCall(responseObserverT);
   }
 
   @Override
   public FirestoreOptions getOptions() {
     return firestoreOptions;
   }
+
+  @Override
+  public void close() throws Exception {
+    firestoreClient.close();
+    closed = true;
+  }
 }
diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITSystemTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITSystemTest.java
@@ -62,6 +62,7 @@
 import java.util.concurrent.Semaphore;
 import java.util.concurrent.atomic.AtomicInteger;
 import javax.annotation.Nullable;
+import org.junit.After;
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
@@ -92,6 +93,11 @@ public void before() {
     randomDoc = randomColl.document();
   }
 
+  @After
+  public void after() throws Exception {
+    firestore.close();
+  }
+
   private DocumentReference addDocument(String key, Object value, Object... fields)
       throws Exception {
     DocumentReference documentReference = randomColl.document();
