Skip to content

Commit 339595d

Browse files
mar-kolyamar-kolya
committed
Try to fix a flaky test
1 parent 320289a commit 339595d

2 files changed

Lines changed: 37 additions & 11 deletions

File tree

‎dd-java-agent/agent-profiling/profiling-uploader/src/main/java/com/datadog/profiling/uploader/ProfileUploader.java‎

Lines changed: 14 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -127,21 +127,23 @@ public void onResponse(final Call call, final Response response) {
127127

128128
ioLogger.error("Failed to upload profile", getLoggerResponse(response));
129129
}
130+
// Note: this whole callback never touches body and would be perfectly happy even if server
131+
// never sends it.
130132
response.close();
131133
}
132134

133-
private static IOLogger.Response getLoggerResponse(okhttp3.Response response) {
135+
private static IOLogger.Response getLoggerResponse(final okhttp3.Response response) {
134136
if (response != null) {
135137
try {
136138
return new IOLogger.Response(
137139
response.code(), response.message(), response.body().string().trim());
138-
} catch (NullPointerException | IOException ignored) {
140+
} catch (final NullPointerException | IOException ignored) {
139141
}
140142
}
141143
return null;
142144
}
143145

144-
private static boolean isEmptyReplyFromServer(IOException e) {
146+
private static boolean isEmptyReplyFromServer(final IOException e) {
145147
// The server in datadog-agent triggers 'unexpected end of stream' caused by EOFException.
146148
// The MockWebServer in tests triggers an InterruptedIOException with SocketPolicy
147149
// NO_RESPONSE. This is because in tests we can't cleanly terminate the connection on the
@@ -162,24 +164,30 @@ private static boolean isEmptyReplyFromServer(IOException e) {
162164
private final String apiKey;
163165
private final String url;
164166
private final String containerId;
167+
private final int terminationTimeout;
165168
private final List<String> tags;
166169
private final Compression compression;
167170
private final Deque<Integer> requestSizeHistory;
168171

169172
public ProfileUploader(final Config config) {
170-
this(config, new IOLogger(log), ContainerInfo.get().getContainerId());
173+
this(config, new IOLogger(log), ContainerInfo.get().getContainerId(), TERMINATION_TIMEOUT);
171174
}
172175

173176
/**
174177
* Note that this method is only visible for testing and should not be used from outside this
175178
* class.
176179
*/
177-
ProfileUploader(final Config config, final IOLogger ioLogger, final String containerId) {
180+
ProfileUploader(
181+
final Config config,
182+
final IOLogger ioLogger,
183+
final String containerId,
184+
final int terminationTimeout) {
178185
url = config.getFinalProfilingUrl();
179186
apiKey = config.getApiKey();
180187
agentless = config.isProfilingAgentless();
181188
responseCallback = new ResponseCallback(ioLogger);
182189
this.containerId = containerId;
190+
this.terminationTimeout = terminationTimeout;
183191

184192
log.debug("Started ProfileUploader with target url {}", url);
185193
/*
@@ -287,13 +295,12 @@ public void upload(final RecordingType type, final RecordingData data) {
287295
public void shutdown() {
288296
okHttpExecutorService.shutdownNow();
289297
try {
290-
okHttpExecutorService.awaitTermination(TERMINATION_TIMEOUT, TimeUnit.SECONDS);
298+
okHttpExecutorService.awaitTermination(terminationTimeout, TimeUnit.SECONDS);
291299
} catch (final InterruptedException e) {
292300
// Note: this should only happen in main thread right before exiting, so eating up interrupted
293301
// state should be fine.
294302
log.warn("Wait for executor shutdown interrupted");
295303
}
296-
297304
client.connectionPool().evictAll();
298305
}
299306

‎dd-java-agent/agent-profiling/profiling-uploader/src/test/java/com/datadog/profiling/uploader/ProfileUploaderTest.java‎

Lines changed: 23 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,7 @@
6161
import okhttp3.mockwebserver.MockResponse;
6262
import okhttp3.mockwebserver.MockWebServer;
6363
import okhttp3.mockwebserver.RecordedRequest;
64+
import okhttp3.mockwebserver.SocketPolicy;
6465
import org.junit.jupiter.api.AfterEach;
6566
import org.junit.jupiter.api.BeforeEach;
6667
import org.junit.jupiter.api.Test;
@@ -112,6 +113,10 @@ public class ProfileUploaderTest {
112113
private final Duration REQUEST_TIMEOUT = Duration.ofSeconds(10);
113114
private final Duration REQUEST_IO_OPERATION_TIMEOUT = Duration.ofSeconds(5);
114115

116+
// Termination timeout has to be longer than request timeout to make sure that all callbacks are
117+
// called before the termination.
118+
private final Duration TERMINATION_TIMEOUT = REQUEST_TIMEOUT.plus(Duration.ofSeconds(5));
119+
115120
private final Duration FOREVER_REQUEST_TIMEOUT = Duration.ofSeconds(1000);
116121

117122
@Mock private Config config;
@@ -133,7 +138,9 @@ public void setup() throws IOException {
133138
when(config.getMergedProfilingTags()).thenReturn(TAGS);
134139
when(config.getProfilingUploadTimeout()).thenReturn((int) REQUEST_TIMEOUT.getSeconds());
135140

136-
uploader = new ProfileUploader(config, ioLogger, "containerId");
141+
uploader =
142+
new ProfileUploader(
143+
config, ioLogger, "containerId", (int) TERMINATION_TIMEOUT.getSeconds());
137144
}
138145

139146
@AfterEach
@@ -202,7 +209,9 @@ public void testRequestParameters(final String compression)
202209

203210
@Test
204211
public void testRequestWithContainerId() throws IOException, InterruptedException {
205-
uploader = new ProfileUploader(config, ioLogger, "container-id");
212+
uploader =
213+
new ProfileUploader(
214+
config, ioLogger, "container-id", (int) TERMINATION_TIMEOUT.getSeconds());
206215

207216
server.enqueue(new MockResponse().setResponseCode(200));
208217
uploader.upload(RECORDING_TYPE, mockRecordingData(RECORDING_RESOURCE));
@@ -390,10 +399,13 @@ public void testConnectionRefused() throws IOException {
390399

391400
@Test
392401
public void testNoReplyFromServer() throws IOException, InterruptedException {
393-
server.enqueue(new MockResponse().setBodyDelay(10, TimeUnit.SECONDS));
402+
server.enqueue(new MockResponse().setSocketPolicy(SocketPolicy.NO_RESPONSE));
394403
final RecordingData recording = mockRecordingData(RECORDING_RESOURCE);
395404
uploader.upload(RECORDING_TYPE, recording);
396405

406+
// Wait longer than request timeout
407+
assertNotNull(server.takeRequest(REQUEST_TIMEOUT.getSeconds() + 1, TimeUnit.SECONDS));
408+
397409
// Shutting down uploader ensures all callbacks are called on http client
398410
uploader.shutdown();
399411
verify(recording.getStream()).close();
@@ -412,10 +424,17 @@ public void testTimeout() throws IOException, InterruptedException {
412424
final RecordingData recording = mockRecordingData(RECORDING_RESOURCE);
413425
uploader.upload(RECORDING_TYPE, recording);
414426

415-
assertNotNull(server.takeRequest(5, TimeUnit.SECONDS));
427+
// Wait longer than request timeout
428+
assertNotNull(
429+
server.takeRequest(REQUEST_IO_OPERATION_TIMEOUT.getSeconds() + 2, TimeUnit.SECONDS));
416430

431+
// Shutting down uploader ensures all callbacks are called on http client
432+
uploader.shutdown();
417433
verify(recording.getStream()).close();
418434
verify(recording).release();
435+
// This seems to be a weird behaviour on okHttp side: it considers request to be a success even
436+
// if it didn't get headers before the timeout
437+
verify(ioLogger).success(eq("Upload done"));
419438
}
420439

421440
@Test

0 commit comments

Comments
 (0)