Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,9 @@
import org.opensearch.dataprepper.plugins.source.microsoft_office365.auth.Office365AuthenticationInterface;
import org.opensearch.dataprepper.plugins.source.source_crawler.exception.SaaSCrawlerException;
import org.opensearch.dataprepper.plugins.source.microsoft_office365.models.AuditLogsResponse;
import org.opensearch.dataprepper.plugins.source.source_crawler.utils.retry.RetryHandler;
import org.opensearch.dataprepper.plugins.source.source_crawler.utils.retry.DefaultRetryStrategy;
import org.opensearch.dataprepper.plugins.source.source_crawler.utils.retry.DefaultStatusCodeHandler;
import org.springframework.core.ParameterizedTypeReference;
import org.springframework.http.HttpEntity;
import org.springframework.http.HttpHeaders;
Expand All @@ -30,6 +33,7 @@
import java.time.Instant;
import java.util.List;
import java.util.Map;
import java.util.Optional;

import static org.opensearch.dataprepper.logging.DataPrepperMarkers.NOISY;
import static org.opensearch.dataprepper.plugins.source.microsoft_office365.utils.Constants.CONTENT_TYPES;
Expand Down Expand Up @@ -57,6 +61,7 @@ public class Office365RestClient {
private static final String MANAGEMENT_API_BASE_URL = "https://manage.office.com/api/v1.0/";

private final RestTemplate restTemplate = new RestTemplate();
private final RetryHandler retryHandler;
private final Office365AuthenticationInterface authConfig;
private final Timer auditLogFetchLatencyTimer;
private final Timer searchCallLatencyTimer;
Expand All @@ -76,6 +81,9 @@ public Office365RestClient(final Office365AuthenticationInterface authConfig,
this.auditLogsRequestedCounter = pluginMetrics.counter(AUDIT_LOGS_REQUESTED);
this.apiCallsCounter = pluginMetrics.counter(API_CALLS);
this.errorTypeMetricCounterMap = getErrorTypeMetricCounterMap(pluginMetrics);
this.retryHandler = new RetryHandler(
new DefaultRetryStrategy(),
new DefaultStatusCodeHandler());

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What do you think if we pass "pluginMetrics" into RetryHandler here so we do not have to pass Optional.of(auditLogRequestsFailedCounter) on line #261? Can also be an Optional param.

IMO, it would keep the executeWithRetry clean and concise.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We don't think it would be beneficial. Even if we pass pluginMetrics into the RetryHandler constructor, we would still need to provide the metric name at line 261. Additionally, inside RetryHandler we would have to iterate through pluginMetrics to locate that specific metric. Please correct us if we are wrong.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You're right. I was thinking of passing this searchRequestsFailedCounter directly in the constructor. E.g

 new RetryHandler(
        new DefaultRetryStrategy(),
        new DefaultStatusCodeHandler(),
        Optional<Counter> failedCounter
);

Not a big blocker and will let others chime in.

}

/**
Expand Down Expand Up @@ -111,7 +119,7 @@ public void startSubscriptions() {
authConfig.getTenantId(),
contentType);

RetryHandler.executeWithRetry(() -> {
retryHandler.executeWithRetry(() -> {
try {
headers.setBearerAuth(authConfig.getAccessToken());
apiCallsCounter.increment();
Expand Down Expand Up @@ -168,7 +176,7 @@ public AuditLogsResponse searchAuditLogs(final String contentType,

return searchCallLatencyTimer.record(() -> {
try {
return RetryHandler.executeWithRetry(
return retryHandler.executeWithRetry(
() -> {
headers.setBearerAuth(authConfig.getAccessToken());
apiCallsCounter.increment();
Expand Down Expand Up @@ -213,7 +221,7 @@ public AuditLogsResponse searchAuditLogs(final String contentType,
return new AuditLogsResponse(response.getBody(), nextPageUri);
},
authConfig::renewCredentials,
provideSearchRequestFailureCounter(pluginMetrics)
Optional.of(provideSearchRequestFailureCounter(pluginMetrics))
);
} catch (Exception e) {
publishErrorTypeMetricCounter(e, this.errorTypeMetricCounterMap);
Expand Down Expand Up @@ -242,7 +250,7 @@ public String getAuditLog(String contentUri) {

return auditLogFetchLatencyTimer.record(() -> {
try {
String response = RetryHandler.executeWithRetry(() -> {
String response = retryHandler.executeWithRetry(() -> {
headers.setBearerAuth(authConfig.getAccessToken());
apiCallsCounter.increment();
ResponseEntity<String> responseEntity = restTemplate.exchange(
Expand All @@ -253,7 +261,7 @@ public String getAuditLog(String contentUri) {
);

return responseEntity.getBody();
}, authConfig::renewCredentials, provideGetRequestsFailureCounter(pluginMetrics));
}, authConfig::renewCredentials, Optional.of(provideGetRequestsFailureCounter(pluginMetrics)));

// Log response details
if (response == null) {
Expand Down

This file was deleted.

Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,9 @@

import lombok.Getter;
import org.opensearch.dataprepper.plugins.source.microsoft_office365.Office365SourceConfig;
import org.opensearch.dataprepper.plugins.source.microsoft_office365.RetryHandler;
import org.opensearch.dataprepper.plugins.source.source_crawler.utils.retry.RetryHandler;
import org.opensearch.dataprepper.plugins.source.source_crawler.utils.retry.DefaultRetryStrategy;
import org.opensearch.dataprepper.plugins.source.source_crawler.utils.retry.DefaultStatusCodeHandler;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.http.HttpEntity;
Expand Down Expand Up @@ -40,6 +42,7 @@ public class Office365AuthenticationProvider implements Office365AuthenticationI
"&scope=%s";

private final RestTemplate restTemplate = new RestTemplate();
private final RetryHandler retryHandler;
private final String tenantId;
private final Office365SourceConfig office365SourceConfig;
private String accessToken;
Expand All @@ -53,6 +56,9 @@ public class Office365AuthenticationProvider implements Office365AuthenticationI
public Office365AuthenticationProvider(Office365SourceConfig config) {
this.tenantId = config.getTenantId();
this.office365SourceConfig = config;
this.retryHandler = new RetryHandler(
new DefaultRetryStrategy(),
new DefaultStatusCodeHandler());
}

@Override
Expand All @@ -76,7 +82,7 @@ public void renewCredentials() {
HttpEntity<String> entity = new HttpEntity<>(payload, headers);
String tokenEndpoint = String.format(TOKEN_URL, office365SourceConfig.getTenantId());

ResponseEntity<Map> response = RetryHandler.executeWithRetry(
ResponseEntity<Map> response = retryHandler.executeWithRetry(
() -> restTemplate.postForEntity(tokenEndpoint, entity, Map.class),
() -> {
} // No credential renewal for authentication endpoint
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,95 @@
/*
* Copyright OpenSearch Contributors
* SPDX-License-Identifier: Apache-2.0
*
* The OpenSearch Contributors require contributions made to
* this file be licensed under the Apache-2.0 license or a
* compatible open source license.
*
*/

package org.opensearch.dataprepper.plugins.source.source_crawler.utils.retry;

import lombok.extern.slf4j.Slf4j;
import org.springframework.http.HttpStatus;

import java.util.Arrays;
import java.util.List;
import java.util.Optional;

/**
* Default retry strategy with fixed backoff times
*/
@Slf4j
public class DefaultRetryStrategy implements RetryStrategy {
private static final List<HttpStatus> DEFAULT_RATE_LIMIT_STATUS_CODES = Arrays.asList(HttpStatus.TOO_MANY_REQUESTS);

private final List<Integer> retryAttemptSleepTime;
private final List<Integer> rateLimitRetrySleepTime;
private final List<HttpStatus> rateLimitStatusCodes;
private final int maxRetries;

/**
* Constructor with default sleep times
*/
public DefaultRetryStrategy() {
this.retryAttemptSleepTime = RetryStrategy.DEFAULT_RETRY_ATTEMPT_SLEEP_TIME;
this.rateLimitRetrySleepTime = RetryStrategy.DEFAULT_RATE_LIMIT_RETRY_SLEEP_TIME;
this.rateLimitStatusCodes = DEFAULT_RATE_LIMIT_STATUS_CODES;
this.maxRetries = RetryStrategy.MAX_RETRIES;
}

/**
* Constructor with custom max retries
*
* @param maxRetries Maximum number of retries
*/
public DefaultRetryStrategy(final int maxRetries) {
this.retryAttemptSleepTime = RetryStrategy.DEFAULT_RETRY_ATTEMPT_SLEEP_TIME;
this.rateLimitRetrySleepTime = RetryStrategy.DEFAULT_RATE_LIMIT_RETRY_SLEEP_TIME;
this.rateLimitStatusCodes = DEFAULT_RATE_LIMIT_STATUS_CODES;
this.maxRetries = maxRetries;
}

/**
* Constructor with Custom sleep times for rate limit retries and custom rate limit status codes
*
* @param rateLimitRetrySleepTime Custom sleep times for rate limit retries (in
* seconds)
* @param rateLimitStatusCodes List of status codes that are considered rate limited
*/
public DefaultRetryStrategy(List<Integer> rateLimitRetrySleepTime, List<HttpStatus> rateLimitStatusCodes) {
this.retryAttemptSleepTime = RetryStrategy.DEFAULT_RETRY_ATTEMPT_SLEEP_TIME;
this.rateLimitRetrySleepTime = rateLimitRetrySleepTime != null
? rateLimitRetrySleepTime
: RetryStrategy.DEFAULT_RATE_LIMIT_RETRY_SLEEP_TIME;
this.rateLimitStatusCodes = rateLimitStatusCodes != null
? rateLimitStatusCodes
: DEFAULT_RATE_LIMIT_STATUS_CODES;
this.maxRetries = this.rateLimitRetrySleepTime.size();
}

@Override
public long calculateSleepTime(Exception ex, int retryCount) {
Optional<HttpStatus> statusCode = RetryStrategy.getStatusCode(ex);

List<Integer> sleepTimes = (statusCode.isPresent() && rateLimitStatusCodes.contains(statusCode.get()))
? rateLimitRetrySleepTime
: retryAttemptSleepTime;

int sleepTimeSeconds = (retryCount < sleepTimes.size())
? sleepTimes.get(retryCount)
: sleepTimes.get(sleepTimes.size() - 1);

log.debug("Retrying in {} seconds (attempt {}/{})",
sleepTimeSeconds, retryCount + 1, getMaxRetries());

return sleepTimeSeconds * RetryStrategy.SLEEP_TIME_MULTIPLIER_MS;
}

@Override
public int getMaxRetries() {
return maxRetries;
}

}
Original file line number Diff line number Diff line change
@@ -0,0 +1,73 @@
/*
* Copyright OpenSearch Contributors
* SPDX-License-Identifier: Apache-2.0
*
* The OpenSearch Contributors require contributions made to
* this file be licensed under the Apache-2.0 license or a
* compatible open source license.
*
*/

package org.opensearch.dataprepper.plugins.source.source_crawler.utils.retry;

import lombok.extern.slf4j.Slf4j;
import org.springframework.http.HttpStatus;
import java.util.Optional;

import static org.opensearch.dataprepper.logging.DataPrepperMarkers.NOISY;

/**
* Default status code handling - covers common HTTP scenarios
*/
@Slf4j
public class DefaultStatusCodeHandler implements StatusCodeHandler {

@vecheka vecheka Nov 25, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We recently found an issue in WindowEvents connector where "tokenExpired" throws 403(FORBIDDEN) instead of 401(UNAUTHORIZED), and we did not run credentialsRenew on it. This caused the connector to stop retrieving data after.

Should we consider allowing connectors to pass information whether to retry credsRenewal on UNAUTHORIZED, FORBIDDEN or both? E.g

Map<HttpStatus, Boolean> authStatusCredentialRenewalMap = new HashMap<>();
authStatusCredentialRenewalMap.put(FORBIDDEN, true);
new DefaultStatusCodeHandler(authStatusCredentialRenewalMap);


.......


case FORBIDDEN:
      log.error(NOISY, "Access forbidden: {}", statusMessage, ex);
      if (authStatusCredentialRenewalMap.get(FORBIDDEN)) {
           credentialRenewal.run();
           return RetryDecision.retry();
       }
      return RetryDecision.stopWithException(
              new SecurityException("Access forbidden: " + statusMessage));

Just one idea, please let me know what you think.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, for cases like this, we have a CustomApiStatusCodeHandler where we can add all status codes that require special handling or are not covered inside DefaultStatusCodeHandler. For all other cases, we simply call super.handleStatusCode as shown below:

class CustomApiStatusCodeHandler extends DefaultStatusCodeHandler {

    @Override
    public RetryDecision handleStatusCode(Exception ex, int retryCount,
                                          Runnable credentialRenewal) {
        Optional<HttpStatus> statusCode = RetryStrategy.getStatusCode(ex);
        String statusMessage = ex.getMessage();

        if (statusCode.isEmpty()) {
            return RetryDecision.stop();
        }

        switch (statusCode.get()) {
            case FORBIDDEN:
                credentialRenewal.run();
                return RetryDecision.retry();

            default:
                // Delegate to default handler for all standard status codes
                return super.handleStatusCode(ex, retryCount, credentialRenewal);
        }
    }
}

And in the RestClient:

this.retryHandler = new RetryHandler(
        new RetryAfterHeaderStrategy(),
        new CustomApiStatusCodeHandler()
);

We have covered all such scenarios in the design document, please check: https://quip-amazon.com/vWbpATaYxCeS/Design-And-Refactor-Retry-Handler-To-Move-Into-Source-Crawler-Package.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Okay makes sense, thanks! That's super helpful.


@Override
public RetryDecision handleStatusCode(Exception ex, int retryCount,
Runnable credentialRenewal) {
Optional<HttpStatus> statusCode = RetryStrategy.getStatusCode(ex);
String statusMessage = ex.getMessage();

if (statusCode.isEmpty()) {
return RetryDecision.stop();
}

switch (statusCode.get()) {
case UNAUTHORIZED:
log.error(NOISY, "Token expired. Attempting to renew credentials.", ex);

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is it possible to have an authFailures metric incremented here? We need it to map to some of ATP's metrics and I believe we don't emit it anywhere. Source

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We're already doing it here.

This is the metricMapping ->

errorTypeMetricCounterMap.put(HttpStatus.UNAUTHORIZED.getReasonPhrase(), pluginMetrics.counter(REQUEST_ACCESS_DENIED));
. Please correct me if you're referring to something else.

credentialRenewal.run();
return RetryDecision.retry();

case FORBIDDEN:
log.error(NOISY, "Access forbidden: {}", statusMessage, ex);
return RetryDecision.stopWithException(
new SecurityException("Access forbidden: " + statusMessage));

case NOT_FOUND:
log.warn(NOISY, "Resource not found (404): {}. " +
"This is expected for deleted/expired resources.", statusMessage);
return RetryDecision.stop();

case TOO_MANY_REQUESTS:
log.error(NOISY, "Hitting API rate limit. Backing off.", ex);
return RetryDecision.retry();

case SERVICE_UNAVAILABLE:
log.error(NOISY, "Service unavailable. Will retry.", ex);
return RetryDecision.retry();

default:
if (statusCode.get().is4xxClientError()) {
log.error(NOISY, "Client error: {}. Will not retry.", statusCode, ex);
return RetryDecision.stop();
} else if (statusCode.get().is5xxServerError()) {
log.error(NOISY, "Server error: {}. Will retry.", statusCode, ex);
return RetryDecision.retry();
} else {
log.error(NOISY, "Unexpected status code: {}. Will not retry.",
statusCode, ex);
return RetryDecision.stop();
}
}
}
}
Loading
Loading