1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
|
/*
* Licensed to Elasticsearch under one or more contributor
* license agreements. See the NOTICE file distributed with
* this work for additional information regarding copyright
* ownership. Elasticsearch licenses this file to you under
* the Apache License, Version 2.0 (the "License"); you may
* not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
package org.elasticsearch.discovery.gce;
import com.google.api.client.auth.oauth2.Credential;
import com.google.api.client.googleapis.testing.auth.oauth2.MockGoogleCredential;
import com.google.api.client.http.*;
import com.google.api.client.util.ExponentialBackOff;
import com.google.api.client.util.Sleeper;
import org.elasticsearch.SpecialPermission;
import org.elasticsearch.common.logging.ESLogger;
import org.elasticsearch.common.logging.ESLoggerFactory;
import java.io.IOException;
import java.security.AccessController;
import java.security.PrivilegedAction;
import java.util.Objects;
public class RetryHttpInitializerWrapper implements HttpRequestInitializer {
private int maxWait;
private static final ESLogger logger =
ESLoggerFactory.getLogger(RetryHttpInitializerWrapper.class.getName());
// Intercepts the request for filling in the "Authorization"
// header field, as well as recovering from certain unsuccessful
// error codes wherein the Credential must refresh its token for a
// retry.
private final Credential wrappedCredential;
// A sleeper; you can replace it with a mock in your test.
private final Sleeper sleeper;
public RetryHttpInitializerWrapper(Credential wrappedCredential) {
this(wrappedCredential, Sleeper.DEFAULT, ExponentialBackOff.DEFAULT_MAX_ELAPSED_TIME_MILLIS);
}
public RetryHttpInitializerWrapper(Credential wrappedCredential, int maxWait) {
this(wrappedCredential, Sleeper.DEFAULT, maxWait);
}
// Use only for testing.
RetryHttpInitializerWrapper(
Credential wrappedCredential, Sleeper sleeper, int maxWait) {
this.wrappedCredential = Objects.requireNonNull(wrappedCredential);
this.sleeper = sleeper;
this.maxWait = maxWait;
}
// Use only for testing
static MockGoogleCredential.Builder newMockCredentialBuilder() {
// TODO: figure out why GCE is so bad like this
SecurityManager sm = System.getSecurityManager();
if (sm != null) {
sm.checkPermission(new SpecialPermission());
}
return AccessController.doPrivileged(new PrivilegedAction<MockGoogleCredential.Builder>() {
@Override
public MockGoogleCredential.Builder run() {
return new MockGoogleCredential.Builder();
}
});
}
@Override
public void initialize(HttpRequest httpRequest) {
final HttpUnsuccessfulResponseHandler backoffHandler =
new HttpBackOffUnsuccessfulResponseHandler(
new ExponentialBackOff.Builder()
.setMaxElapsedTimeMillis(maxWait)
.build())
.setSleeper(sleeper);
httpRequest.setInterceptor(wrappedCredential);
httpRequest.setUnsuccessfulResponseHandler(
new HttpUnsuccessfulResponseHandler() {
int retry = 0;
@Override
public boolean handleResponse(HttpRequest request, HttpResponse response, boolean supportsRetry) throws IOException {
if (wrappedCredential.handleResponse(
request, response, supportsRetry)) {
// If credential decides it can handle it,
// the return code or message indicated
// something specific to authentication,
// and no backoff is desired.
return true;
} else if (backoffHandler.handleResponse(
request, response, supportsRetry)) {
// Otherwise, we defer to the judgement of
// our internal backoff handler.
logger.debug("Retrying [{}] times : [{}]", retry, request.getUrl());
return true;
} else {
return false;
}
}
});
httpRequest.setIOExceptionHandler(
new HttpBackOffIOExceptionHandler(
new ExponentialBackOff.Builder()
.setMaxElapsedTimeMillis(maxWait)
.build())
.setSleeper(sleeper)
);
}
}
|
