logging: make flush wait for writes

google-cloud-logging/src/main/java/com/google/cloud/logging/LoggingHandler.java

@@ -31,8 +31,6 @@
 import java.util.IdentityHashMap;
 import java.util.List;
 import java.util.Set;
-import java.util.concurrent.locks.Lock;
-import java.util.concurrent.locks.ReentrantLock;
 import java.util.logging.ErrorManager;
 import java.util.logging.Filter;
 import java.util.logging.Formatter;
@@ -125,7 +123,7 @@ public class LoggingHandler extends Handler {
   // https://github.com/GoogleCloudPlatform/google-cloud-java/issues/1740 .
   private final Level baseLevel;
 
-  private final Lock writeLock = new ReentrantLock();
+  private final Object writeLock = new Object();
   private final Set<ApiFuture<Void>> pendingWrites =
       Collections.newSetFromMap(new IdentityHashMap<ApiFuture<Void>, Boolean>());
 
@@ -473,25 +471,23 @@ void write(LogEntry entry, WriteOption... options) {
       case ASYNC:
       default:
         final ApiFuture<Void> writeFuture = getLogging().writeAsync(entryList, options);
-        writeLock.lock();
-        try {
+        synchronized(writeLock) {
           pendingWrites.add(writeFuture);
-        } finally {
-          writeLock.unlock();
         }
         ApiFutures.addCallback(
             writeFuture,
             new ApiFutureCallback<Void>() {
-              @Override
-              public void onSuccess(Void v) {
-                writeLock.lock();
-                try {
+              private void removeFromPending() {
+                synchronized(writeLock) {
                   pendingWrites.remove(writeFuture);
-                } finally {
-                  writeLock.unlock();
                 }
               }
 
+              @Override
+              public void onSuccess(Void v) {
+                removeFromPending();
+              }
+
               @Override
               public void onFailure(Throwable t) {
                 try {
@@ -501,7 +497,7 @@ public void onFailure(Throwable t) {
                     reportError(null, new Exception(t), ErrorManager.FLUSH_FAILURE);
                   }
                 } finally {
-                  onSuccess(null);
+                  removeFromPending();
                 }
               }
             });
@@ -511,16 +507,14 @@ public void onFailure(Throwable t) {
 
   @Override
   public void flush() {
-    // BUG(1795): flush is broken, need support from batching implementation.
+    // BUG(1795): We should force batcher to issue RPC call for buffered messages,
+    // so the code below doesn't wait uselessly.
 
-    ArrayList<ApiFuture<Void>> writes = new ArrayList<>();
-    writeLock.lock();
-    try {
-      writes.addAll(pendingWrites);
-    } finally {
-      writeLock.unlock();
+    ArrayList<ApiFuture<Void>> writesToFlush = new ArrayList<>();
+    synchronized(writeLock) {
+      writesToFlush.addAll(pendingWrites);
     }
-    for (ApiFuture<Void> write : writes) {
+    for (ApiFuture<Void> write : writesToFlush) {
       try {
         Uninterruptibles.getUninterruptibly(write);
       } catch (Exception e) {

google-cloud-logging/src/test/java/com/google/cloud/logging/LoggingHandlerTest.java

@@ -18,8 +18,10 @@
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
 
 import com.google.api.gax.core.ApiFutures;
+import com.google.api.gax.core.SettableApiFuture;
 import com.google.cloud.MonitoredResource;
 import com.google.cloud.logging.LogEntry.Builder;
 import com.google.cloud.logging.Logging.WriteOption;
@@ -380,6 +382,40 @@ public void testFlushLevel() {
     handler.publish(newLogRecord(Level.WARNING, MESSAGE));
   }
 
+  @Test
+  public void testFlush() throws InterruptedException {
+    final SettableApiFuture<Void> mockRpc = SettableApiFuture.create();
+
+    EasyMock.expect(options.getProjectId()).andReturn(PROJECT).anyTimes();
+    EasyMock.expect(options.getService()).andReturn(logging);
+    logging.writeAsync(ImmutableList.of(INFO_ENTRY), DEFAULT_OPTIONS);
+    EasyMock.expectLastCall().andReturn(mockRpc);
+    EasyMock.replay(options, logging);
+    final LoggingHandler handler = new LoggingHandler(LOG_NAME, options);
+    handler.setFormatter(new TestFormatter());
+
+    // no messages, nothing to flush.
+    handler.flush();
+
+    // send a message
+    handler.publish(newLogRecord(Level.INFO, MESSAGE));
+    Thread flushWaiter = new Thread(new Runnable() {
+      @Override
+      public void run() {
+        handler.flush();
+      }
+    });
+    flushWaiter.start();
+
+    // flushWaiter should be waiting for mockRpc to complete.
+    flushWaiter.join(100);
+    assertTrue(flushWaiter.isAlive());
+
+    // With the RPC completed, flush should return, and the thread should terminate.
+    mockRpc.set(null);
+    flushWaiter.join();
+  }
+
   @Test
   public void testSyncWrite() {
     EasyMock.expect(options.getProjectId()).andReturn(PROJECT).anyTimes();