Increment Office365 request failure metrics on retry (#6204)

Signed-off-by: Brendan Benner <bbenner@amazon.com>
Co-authored-by: Brendan Benner <bbenner@amazon.com>

Author: bbenner7635

data-prepper-plugins/saas-source-plugins/microsoft-office365-source/src/main/java/org/opensearch/dataprepper/plugins/source/microsoft_office365/Office365RestClient.java

@@ -182,11 +182,11 @@ public AuditLogsResponse searchAuditLogs(final String contentType,
 
                             return new AuditLogsResponse(response.getBody(), nextPageUri);
                         },
-                        authConfig::renewCredentials
+                        authConfig::renewCredentials,
+                        searchRequestsFailedCounter
                 );
             } catch (Exception e) {
                 log.error(NOISY, "Error while fetching audit logs for content type {}", contentType, e);
-                searchRequestsFailedCounter.increment();
                 throw new RuntimeException("Failed to fetch audit logs", e);
             }
         });
@@ -216,12 +216,11 @@ public String getAuditLog(String contentUri) {
                             new HttpEntity<>(headers),
                             String.class
                     ).getBody();
-                }, authConfig::renewCredentials);
+                }, authConfig::renewCredentials, auditLogRequestsFailedCounter);
                 auditLogRequestsSuccessCounter.increment();
                 return response;
             } catch (Exception e) {
                 log.error(NOISY, "Error while fetching audit log content from URI: {}", contentUri, e);
-                auditLogRequestsFailedCounter.increment();
                 throw new RuntimeException("Failed to fetch audit log", e);
             }
         });

data-prepper-plugins/saas-source-plugins/microsoft-office365-source/src/main/java/org/opensearch/dataprepper/plugins/source/microsoft_office365/RetryHandler.java

@@ -9,6 +9,7 @@
 
 package org.opensearch.dataprepper.plugins.source.microsoft_office365;
 
+import io.micrometer.core.instrument.Counter;
 import lombok.extern.slf4j.Slf4j;
 import org.springframework.http.HttpStatus;
 import org.springframework.web.client.HttpClientErrorException;
@@ -26,10 +27,17 @@ public class RetryHandler {
     private static final int SLEEP_TIME_MULTIPLIER = 1000;
 
     public static <T> T executeWithRetry(Supplier<T> operation, Runnable credentialRenewal) {
+        return executeWithRetry(operation, credentialRenewal, null);
+    }
+
+    public static <T> T executeWithRetry(Supplier<T> operation, Runnable credentialRenewal, Counter failureCounter) {
         int retryCount = 0;
         while (retryCount < MAX_RETRIES) {
+            boolean operationSucceeded = false;
             try {
-                return operation.get();
+                T result = operation.get();
+                operationSucceeded = true;
+                return result;
             } catch (HttpClientErrorException | HttpServerErrorException ex) {
                 HttpStatus statusCode = ex.getStatusCode();
                 String statusMessage = ex.getMessage();
@@ -70,9 +78,13 @@ public static <T> T executeWithRetry(Supplier<T> operation, Runnable credentialR
                     Thread.currentThread().interrupt();
                     throw new RuntimeException("Retry interrupted", ie);
                 }
+            } finally {
+                if (!operationSucceeded && failureCounter != null) {
+                    failureCounter.increment();
+                }
             }
             retryCount++;
         }
         throw new RuntimeException("Exceeded max retry attempts");
     }
-}
+}

data-prepper-plugins/saas-source-plugins/microsoft-office365-source/src/test/java/org/opensearch/dataprepper/plugins/source/microsoft_office365/Office365RestClientTest.java

@@ -12,6 +12,7 @@
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
+import io.micrometer.core.instrument.Counter;
 import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
 import org.mockito.junit.jupiter.MockitoExtension;
@@ -291,4 +292,59 @@ void testTokenRenewal() {
         assertEquals("Bearer token-0", requestTokens.get(0), "First request should use token-0");
         assertEquals("Bearer token-1", requestTokens.get(1), "Second request should use token-1");
     }
+
+    @Test
+    void testSearchAuditLogsFailureCounterIncrementsOnEachRetry() throws Exception {
+        Instant startTime = Instant.now().minus(1, ChronoUnit.HOURS);
+        Instant endTime = Instant.now();
+        
+        when(authConfig.getTenantId()).thenReturn("test-tenant-id");
+        when(authConfig.getAccessToken()).thenReturn("test-access-token");
+
+        Counter mockCounter = org.mockito.Mockito.mock(Counter.class);
+        ReflectivelySetField.setField(Office365RestClient.class, office365RestClient, "searchRequestsFailedCounter", mockCounter);
+
+        when(restTemplate.exchange(
+                anyString(),
+                eq(HttpMethod.GET),
+                any(),
+                any(ParameterizedTypeReference.class)
+        )).thenThrow(new HttpClientErrorException(HttpStatus.TOO_MANY_REQUESTS));
+
+        assertThrows(RuntimeException.class, () -> 
+            office365RestClient.searchAuditLogs(
+                    "Audit.AzureActiveDirectory",
+                    startTime,
+                    endTime,
+                    null
+            )
+        );
+
+        // Verify counter.increment() was called exactly 6 times (once for each retry attempt)
+        verify(mockCounter, times(6)).increment();
+    }
+
+    @Test
+    void testGetAuditLogFailureCounterIncrementsOnEachRetry() throws Exception {
+        String contentUri = "https://manage.office.com/api/v1.0/test-tenant/activity/feed/audit/123";
+        
+        when(authConfig.getAccessToken()).thenReturn("test-access-token");
+
+        Counter mockCounter = org.mockito.Mockito.mock(Counter.class);
+        ReflectivelySetField.setField(Office365RestClient.class, office365RestClient, "auditLogRequestsFailedCounter", mockCounter);
+
+        when(restTemplate.exchange(
+                eq(contentUri),
+                eq(HttpMethod.GET),
+                any(),
+                eq(String.class)
+        )).thenThrow(new HttpClientErrorException(HttpStatus.TOO_MANY_REQUESTS));
+
+        assertThrows(RuntimeException.class, () -> 
+            office365RestClient.getAuditLog(contentUri)
+        );
+
+        // Verify counter.increment() was called exactly 6 times (once for each retry attempt)
+        verify(mockCounter, times(6)).increment();
+    }
 }