-
Notifications
You must be signed in to change notification settings - Fork 15
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
1 parent
6a10243
commit 6dc9052
Showing
4 changed files
with
265 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
141 changes: 141 additions & 0 deletions
141
dialogue-core/src/main/java/com/palantir/dialogue/core/RequestSizeMetricsChannel.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,141 @@ | ||
/* | ||
* (c) Copyright 2023 Palantir Technologies Inc. All rights reserved. | ||
* | ||
* Licensed 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 com.palantir.dialogue.core; | ||
|
||
import com.codahale.metrics.Histogram; | ||
import com.google.common.util.concurrent.ListenableFuture; | ||
import com.palantir.conjure.java.client.config.ClientConfiguration; | ||
import com.palantir.dialogue.Endpoint; | ||
import com.palantir.dialogue.EndpointChannel; | ||
import com.palantir.dialogue.Request; | ||
import com.palantir.dialogue.RequestBody; | ||
import com.palantir.dialogue.Response; | ||
import com.palantir.tritium.metrics.registry.TaggedMetricRegistry; | ||
import java.io.FilterOutputStream; | ||
import java.io.IOException; | ||
import java.io.OutputStream; | ||
import java.util.Optional; | ||
|
||
final class RequestSizeMetricsChannel implements EndpointChannel { | ||
private final EndpointChannel delegate; | ||
private final Histogram requestSize; | ||
|
||
static EndpointChannel create(Config cf, EndpointChannel channel, Endpoint endpoint) { | ||
ClientConfiguration clientConf = cf.clientConf(); | ||
return new RequestSizeMetricsChannel(channel, endpoint, clientConf.taggedMetricRegistry()); | ||
} | ||
|
||
RequestSizeMetricsChannel(EndpointChannel delegate, Endpoint endpoint, TaggedMetricRegistry registry) { | ||
this.delegate = delegate; | ||
DialogueClientMetrics dialogueClientMetrics = DialogueClientMetrics.of(registry); | ||
this.requestSize = dialogueClientMetrics | ||
.requestsSize() | ||
.serviceName(endpoint.serviceName()) | ||
.endpoint(endpoint.endpointName()) | ||
.build(); | ||
} | ||
|
||
@Override | ||
public ListenableFuture<Response> execute(Request request) { | ||
Request augmentedRequest = wrap(request); | ||
return delegate.execute(augmentedRequest); | ||
} | ||
|
||
private Request wrap(Request request) { | ||
Optional<RequestBody> body = request.body(); | ||
if (body.isEmpty()) { | ||
// No need to record empty bodies | ||
return request; | ||
} | ||
|
||
return Request.builder() | ||
.from(request) | ||
.body(new RequestSizeRecordingRequestBody(body.get(), this.requestSize)) | ||
.build(); | ||
} | ||
|
||
private class RequestSizeRecordingRequestBody implements RequestBody { | ||
private final RequestBody delegate; | ||
private final Histogram size; | ||
private SizeTrackingOutputStream out; | ||
|
||
RequestSizeRecordingRequestBody(RequestBody requestBody, Histogram size) { | ||
this.delegate = requestBody; | ||
this.size = size; | ||
// we'll never actually write to this output stream, but is safe to perform all operations on in case a | ||
// client closes without calling write. | ||
this.out = new SizeTrackingOutputStream(OutputStream.nullOutputStream(), size); | ||
} | ||
|
||
@Override | ||
public void writeTo(OutputStream output) throws IOException { | ||
out = new SizeTrackingOutputStream(output, size); | ||
delegate.writeTo(out); | ||
} | ||
|
||
@Override | ||
public String contentType() { | ||
return delegate.contentType(); | ||
} | ||
|
||
@Override | ||
public boolean repeatable() { | ||
return delegate.repeatable(); | ||
} | ||
|
||
@Override | ||
public void close() { | ||
try { | ||
out.close(); | ||
} catch (IOException e) { | ||
throw new RuntimeException(e); | ||
} | ||
delegate.close(); | ||
} | ||
|
||
/** | ||
* {@link SizeTrackingOutputStream} records the total number of bytes written to the output stream. | ||
*/ | ||
private final class SizeTrackingOutputStream extends FilterOutputStream { | ||
private final Histogram size; | ||
private long writes = 0; | ||
|
||
SizeTrackingOutputStream(OutputStream delegate, Histogram size) { | ||
super(delegate); | ||
this.size = size; | ||
} | ||
|
||
@Override | ||
public void write(byte[] buffer, int off, int len) throws IOException { | ||
writes += len; | ||
out.write(buffer, off, len); | ||
} | ||
|
||
@Override | ||
public void write(int value) throws IOException { | ||
writes += 1; | ||
out.write(value); | ||
} | ||
|
||
@Override | ||
public void close() throws IOException { | ||
this.size.update(writes); | ||
super.close(); | ||
} | ||
} | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
119 changes: 119 additions & 0 deletions
119
dialogue-core/src/test/java/com/palantir/dialogue/core/RequestSizeMetricsChannelTest.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,119 @@ | ||
/* | ||
* (c) Copyright 2023 Palantir Technologies Inc. All rights reserved. | ||
* | ||
* Licensed 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 com.palantir.dialogue.core; | ||
|
||
import static org.assertj.core.api.Assertions.assertThat; | ||
|
||
import com.codahale.metrics.Snapshot; | ||
import com.google.common.util.concurrent.Futures; | ||
import com.google.common.util.concurrent.ListenableFuture; | ||
import com.palantir.conjure.java.client.config.ClientConfiguration; | ||
import com.palantir.dialogue.EndpointChannel; | ||
import com.palantir.dialogue.Request; | ||
import com.palantir.dialogue.RequestBody; | ||
import com.palantir.dialogue.Response; | ||
import com.palantir.dialogue.TestConfigurations; | ||
import com.palantir.dialogue.TestEndpoint; | ||
import com.palantir.dialogue.TestResponse; | ||
import com.palantir.tritium.metrics.registry.DefaultTaggedMetricRegistry; | ||
import com.palantir.tritium.metrics.registry.TaggedMetricRegistry; | ||
import java.io.ByteArrayOutputStream; | ||
import java.io.IOException; | ||
import java.io.OutputStream; | ||
import java.nio.charset.StandardCharsets; | ||
import java.util.OptionalLong; | ||
import java.util.concurrent.ExecutionException; | ||
import org.junit.jupiter.api.Test; | ||
import org.junit.jupiter.api.extension.ExtendWith; | ||
import org.mockito.Mock; | ||
import org.mockito.junit.jupiter.MockitoExtension; | ||
|
||
@ExtendWith(MockitoExtension.class) | ||
public class RequestSizeMetricsChannelTest { | ||
@Mock | ||
DialogueChannelFactory factory; | ||
|
||
@Test | ||
public void records_request_size_metrics() throws ExecutionException, InterruptedException { | ||
TaggedMetricRegistry registry = DefaultTaggedMetricRegistry.getDefault(); | ||
|
||
ByteArrayOutputStream baos = new ByteArrayOutputStream(); | ||
byte[] expected = "test request body".getBytes(StandardCharsets.UTF_8); | ||
Request request = Request.builder() | ||
.body(new RequestBody() { | ||
@Override | ||
public void writeTo(OutputStream output) throws IOException { | ||
output.write(expected); | ||
} | ||
|
||
@Override | ||
public String contentType() { | ||
return "text/plain"; | ||
} | ||
|
||
@Override | ||
public boolean repeatable() { | ||
return true; | ||
} | ||
|
||
@Override | ||
public OptionalLong contentLength() { | ||
return OptionalLong.of(expected.length); | ||
} | ||
|
||
@Override | ||
public void close() {} | ||
}) | ||
.build(); | ||
|
||
EndpointChannel channel = RequestSizeMetricsChannel.create( | ||
config(ClientConfiguration.builder() | ||
.from(TestConfigurations.create("https://foo:10001")) | ||
.taggedMetricRegistry(registry) | ||
.build()), | ||
r -> { | ||
try { | ||
RequestBody body = r.body().get(); | ||
body.writeTo(baos); | ||
body.close(); | ||
} catch (IOException e) { | ||
throw new RuntimeException(e); | ||
} | ||
return Futures.immediateFuture(new TestResponse().code(200)); | ||
}, | ||
TestEndpoint.GET); | ||
ListenableFuture<Response> response = channel.execute(request); | ||
|
||
assertThat(response.get().code()).isEqualTo(200); | ||
Snapshot snapshot = DialogueClientMetrics.of(registry) | ||
.requestsSize() | ||
.serviceName("service") | ||
.endpoint("endpoint") | ||
.build() | ||
.getSnapshot(); | ||
assertThat(snapshot.size()).isEqualTo(1); | ||
assertThat(snapshot.get99thPercentile()).isEqualTo(expected.length); | ||
} | ||
|
||
private ImmutableConfig config(ClientConfiguration rawConfig) { | ||
return ImmutableConfig.builder() | ||
.channelName("channelName") | ||
.channelFactory(factory) | ||
.rawConfig(rawConfig) | ||
.build(); | ||
} | ||
} |