Address PR comments, fix tests

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

@@ -19,13 +19,16 @@
 import com.google.api.client.util.Strings;
 import com.google.api.gax.core.GaxProperties;
 import com.google.cloud.Tuple;
+import com.google.cloud.logging.Logging.WriteOption;
 import com.google.cloud.logging.Payload.JsonPayload;
 import com.google.cloud.logging.Payload.Type;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Iterables;
 import com.google.protobuf.ListValue;
 import com.google.protobuf.Struct;
 import com.google.protobuf.Value;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.List;
 
 public class Instrumentation {
@@ -38,7 +41,15 @@ public class Instrumentation {
   public static final int MAX_DIAGNOSTIC_VALUE_LENGTH = 14;
   public static final int MAX_DIAGNOSTIC_ENTIES = 3;
   private static boolean instrumentationAdded = false;
+  private static Object instrumentationLock = new Object();
 
+  /**
+   * Populates entries with instrumentation info which is added in separate log entry
+   *
+   * @param logEntries {Iterable<LogEntry>} The list of entries to be populated
+   * @return {Tuple<Boolean, Iterable<LogEntry>>} containg a flag if instrumentation info was added
+   *     or not and a modified list of log entries
+   */
   public static Tuple<Boolean, Iterable<LogEntry>> populateInstrumentationInfo(
       Iterable<LogEntry> logEntries) {
     boolean isWritten = setInstrumentationStatus(true);
@@ -77,6 +88,23 @@ public static Tuple<Boolean, Iterable<LogEntry>> populateInstrumentationInfo(
     return Tuple.of(true, entries);
   }
 
+  /**
+   * Adds a partialSuccess flag option to array of WriteOption
+   *
+   * @param options {WriteOption[]} The options array to be extended
+   * @return The new array of oprions containing WriteOption.OptionType.PARTIAL_SUCCESS flag set to
+   *     true
+   */
+  public static WriteOption[] addPartialSuccessOption(WriteOption[] options) {
+    List<WriteOption> writeOptions = new ArrayList<WriteOption>();
+    writeOptions.addAll(Arrays.asList(options));
+    // Make sure we remove all partial success flags if any exist
+    writeOptions.removeIf(
+        option -> option.getOptionType() == WriteOption.OptionType.PARTIAL_SUCCESS);
+    writeOptions.add(WriteOption.partialSuccess(true));
+    return Iterables.toArray(writeOptions, WriteOption.class);
+  }
+
   /**
    * The helper method to generate a log entry with diagnostic instrumentation data.
    *
@@ -139,7 +167,6 @@ private static ListValue generateLibrariesList(
                 Value.newBuilder().setStructValue(createInfoStruct(name, version)).build());
             if (libraryList.getValuesCount() == MAX_DIAGNOSTIC_ENTIES) break;
           } catch (Exception ex) {
-            System.err.println("ERROR: unexpected exception in generateLibrariesList: " + ex);
           }
         }
       }
@@ -163,11 +190,15 @@ private static Struct createInfoStruct(String libraryName, String libraryVersion
    * already written or not.
    *
    * @param value {boolean} The value to be set.
-   * @returns The value of the flag before it is set.
+   * @returns The value of the flag before it was set.
    */
   public static boolean setInstrumentationStatus(boolean value) {
     if (instrumentationAdded == value) return instrumentationAdded;
-    return setAndGetInstrumentationStatus(value);
+    synchronized (instrumentationLock) {
+      boolean current = instrumentationAdded;
+      instrumentationAdded = value;
+      return current;
+    }
   }
 
   /**
@@ -183,12 +214,6 @@ public static String getLibraryVersion(Class<?> libraryClass) {
     return libraryVersion;
   }
 
-  private static synchronized boolean setAndGetInstrumentationStatus(boolean value) {
-    boolean current = instrumentationAdded;
-    instrumentationAdded = value;
-    return current;
-  }
-
   private static String truncateValue(String value) {
     if (Strings.isNullOrEmpty(value) || value.length() < MAX_DIAGNOSTIC_VALUE_LENGTH) return value;
     return value.substring(0, MAX_DIAGNOSTIC_VALUE_LENGTH) + "*";

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

@@ -19,7 +19,6 @@
 import static com.google.common.base.MoreObjects.firstNonNull;
 
 import com.google.cloud.MonitoredResource;
-import com.google.cloud.Tuple;
 import com.google.cloud.logging.Logging.WriteOption;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
@@ -313,9 +312,10 @@ public void publish(LogRecord record) {
     }
     if (logEntry != null) {
       try {
-        Tuple<Boolean, Iterable<LogEntry>> pair =
-            Instrumentation.populateInstrumentationInfo(ImmutableList.of(logEntry));
-        Iterable<LogEntry> logEntries = pair.y();
+        Iterable<LogEntry> logEntries =
+            redirectToStdout
+                ? Instrumentation.populateInstrumentationInfo(ImmutableList.of(logEntry)).y()
+                : ImmutableList.of(logEntry);
         if (autoPopulateMetadata) {
           logEntries =
               logging.populateMetadata(
@@ -324,15 +324,7 @@ public void publish(LogRecord record) {
         if (redirectToStdout) {
           logEntries.forEach(log -> System.out.println(log.toStructuredJsonString()));
         } else {
-          // Add partialSuccess option always for request containing instrumentation data
-          if (pair.x()) {
-            List<WriteOption> writeOptions = new ArrayList<WriteOption>();
-            writeOptions.addAll(Arrays.asList(defaultWriteOptions));
-            writeOptions.add(WriteOption.partialSuccess(true));
-            logging.write(logEntries, Iterables.toArray(writeOptions, WriteOption.class));
-          } else {
-            logging.write(logEntries, defaultWriteOptions);
-          }
+          logging.write(logEntries, defaultWriteOptions);
         }
       } catch (Exception ex) {
         getErrorManager().error(null, ex, ErrorManager.WRITE_FAILURE);

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

@@ -40,6 +40,7 @@
 import com.google.cloud.MonitoredResource;
 import com.google.cloud.MonitoredResourceDescriptor;
 import com.google.cloud.PageImpl;
+import com.google.cloud.Tuple;
 import com.google.cloud.logging.spi.v2.LoggingRpc;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Function;
@@ -852,15 +853,19 @@ public void write(Iterable<LogEntry> logEntries, WriteOption... options) {
       final Boolean logingOptionsPopulateFlag = getOptions().getAutoPopulateMetadata();
       final Boolean writeOptionPopulateFlga =
           WriteOption.OptionType.AUTO_POPULATE_METADATA.get(writeOptions);
+      Tuple<Boolean, Iterable<LogEntry>> pair =
+          Instrumentation.populateInstrumentationInfo(logEntries);
+      logEntries = pair.y();
 
       if (writeOptionPopulateFlga == Boolean.TRUE
           || (writeOptionPopulateFlga == null && logingOptionsPopulateFlag == Boolean.TRUE)) {
         final MonitoredResource sharedResourceMetadata = RESOURCE.get(writeOptions);
         logEntries =
             populateMetadata(logEntries, sharedResourceMetadata, this.getClass().getName());
       }
-
-      writeLogEntries(logEntries, options);
+      // Add partialSuccess option always for request containing instrumentation data
+      writeLogEntries(
+          logEntries, pair.x() ? Instrumentation.addPartialSuccessOption(options) : options);
       if (flushSeverity != null) {
         for (LogEntry logEntry : logEntries) {
           // flush pending writes if log severity at or above flush severity

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

@@ -629,29 +629,6 @@ public void testRedirectToStdoutDisabled() {
     System.setOut(null);
   }
 
-  @Test
-  public void testDiagnosticInfo() {
-    Instrumentation.setInstrumentationStatus(false);
-    LogEntry json_entry =
-        LogEntry.newBuilder(
-                InstrumentationTest.generateInstrumentationPayload(
-                    Instrumentation.JAVA_LIBRARY_NAME_PREFIX,
-                    Instrumentation.getLibraryVersion(Instrumentation.class.getClass())))
-            .build();
-    logging.write(
-        ImmutableList.of(FINEST_ENTRY, json_entry),
-        WriteOption.logName(LOG_NAME),
-        WriteOption.resource(DEFAULT_RESOURCE),
-        WriteOption.labels(BASE_SEVERITY_MAP),
-        WriteOption.partialSuccess(true));
-    expectLastCall().once();
-    replay(options, logging);
-    LoggingHandler handler = new LoggingHandler(LOG_NAME, options, DEFAULT_RESOURCE);
-    handler.setLevel(Level.ALL);
-    handler.setFormatter(new TestFormatter());
-    handler.publish(newLogRecord(Level.FINEST, MESSAGE));
-  }
-
   private void testPublishCustomResourceWithDestination(
       LogEntry entry, LogDestinationName destination) {
     MonitoredResource resource = MonitoredResource.of("custom", ImmutableMap.<String, String>of());

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

@@ -252,6 +252,7 @@ private void configureListLogsTests(
 
   @Before
   public void setUp() {
+    Instrumentation.setInstrumentationStatus(true);
     rpcFactoryMock = EasyMock.createStrictMock(LoggingRpcFactory.class);
     loggingRpcMock = EasyMock.createStrictMock(LoggingRpc.class);
     EasyMock.expect(rpcFactoryMock.create(EasyMock.anyObject(LoggingOptions.class)))
@@ -2288,6 +2289,41 @@ public void run() {
     assertSame(0, exceptions.get());
   }
 
+  @Test
+  public void testDiagnosticInfoWithNoPartialSuccess() {
+    testDiagnosticInfoGeneration(false);
+  }
+
+  @Test
+  public void testDiagnosticInfoWithPartialSuccess() {
+    testDiagnosticInfoGeneration(true);
+  }
+
+  private void testDiagnosticInfoGeneration(boolean addPartialSuccessOption) {
+    Instrumentation.setInstrumentationStatus(false);
+    LogEntry json_entry =
+        LogEntry.newBuilder(
+                InstrumentationTest.generateInstrumentationPayload(
+                    Instrumentation.JAVA_LIBRARY_NAME_PREFIX,
+                    Instrumentation.getLibraryVersion(Instrumentation.class.getClass())))
+            .build();
+    WriteLogEntriesRequest request =
+        WriteLogEntriesRequest.newBuilder()
+            .addAllEntries(
+                Iterables.transform(
+                    ImmutableList.of(LOG_ENTRY1, LOG_ENTRY2, json_entry),
+                    LogEntry.toPbFunction(PROJECT)))
+            .setPartialSuccess(true)
+            .build();
+    WriteLogEntriesResponse response = WriteLogEntriesResponse.newBuilder().build();
+    EasyMock.expect(loggingRpcMock.write(request)).andReturn(ApiFutures.immediateFuture(response));
+    EasyMock.replay(rpcFactoryMock, loggingRpcMock);
+    logging = options.getService();
+    logging.write(
+        ImmutableList.of(LOG_ENTRY1, LOG_ENTRY2),
+        WriteOption.partialSuccess(addPartialSuccessOption));
+  }
+
   private void testDeleteByDestination(
       String logId, String logName, LogDestinationName destination, boolean useAsyncDelete)
       throws ExecutionException, InterruptedException {