Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Make RequestContextExtension.hook() public API #5107

Merged
merged 33 commits into from
Oct 23, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
33 commits
Select commit Hold shift + click to select a range
ebab297
Make hook() public API for server side
my4-dev Jul 17, 2023
e4e4e3f
save
my4-dev Jul 20, 2023
d2dfd8e
implement hook() public API for client side and move up hook() method
my4-dev Jul 29, 2023
f533660
resolve a build error
my4-dev Jul 30, 2023
f11412b
add tests
my4-dev Aug 1, 2023
15af7cb
revert changes of TransformingResponsePreparationTest
my4-dev Aug 1, 2023
013cedc
fix test failures
my4-dev Aug 1, 2023
2af76dd
merge main branch and resolve conflicts
my4-dev Aug 6, 2023
9cdd0db
add JavaDoc and refactor
my4-dev Aug 6, 2023
e9601ba
add a test and delete comments
my4-dev Aug 6, 2023
197d87f
add tests and revert changes
my4-dev Aug 9, 2023
012603f
apply some comments
my4-dev Aug 11, 2023
cbea40d
change to use NOOP_CONTEXT_HOOK
my4-dev Aug 11, 2023
741ac4e
change contextHook to chain some contextHooks in DefaultServiceConfig…
my4-dev Aug 12, 2023
04be7e7
lint
my4-dev Aug 12, 2023
2a9c7b3
Merge branch 'main' into issue-4991
my4-dev Sep 19, 2023
df37ad7
lint
my4-dev Sep 27, 2023
47bb671
add methods and fix a test case
my4-dev Sep 28, 2023
677b29c
Update core/src/main/java/com/linecorp/armeria/client/ClientOptions.java
ikhoon Oct 13, 2023
c106725
Merge branch 'main' into pr-5107-my4-dev-issue-4991
ikhoon Oct 13, 2023
fc68513
Clean up
ikhoon Oct 13, 2023
f9b0848
merge hooks
ikhoon Oct 13, 2023
da3e278
merge hooks
ikhoon Oct 13, 2023
16e7955
Add integration tests
ikhoon Oct 13, 2023
cc4e315
clean up
ikhoon Oct 13, 2023
716395a
lint
ikhoon Oct 13, 2023
913734f
Merge branch 'main' into pr-5107-my4-dev-issue-4991
ikhoon Oct 18, 2023
00f0a4a
Fix broken tests
ikhoon Oct 18, 2023
9098daa
lint
ikhoon Oct 18, 2023
ecf5921
Merge branch 'main' into pr-5107-my4-dev-issue-4991
minwoox Oct 19, 2023
2738552
Address comments by @trustin
ikhoon Oct 20, 2023
92ff118
Update core/src/main/java/com/linecorp/armeria/client/AbstractClientO…
ikhoon Oct 20, 2023
687bd8b
Update core/src/main/java/com/linecorp/armeria/server/ServiceConfig.java
ikhoon Oct 20, 2023
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,8 @@

package com.linecorp.armeria.server;

import static com.linecorp.armeria.internal.common.RequestContextUtil.NOOP_CONTEXT_HOOK;

import java.nio.file.Path;
import java.util.List;

Expand Down Expand Up @@ -64,19 +66,21 @@ public class RoutersBenchmark {
SERVICE, defaultLogName, defaultServiceName, defaultServiceNaming, 0, 0,
false, AccessLogWriter.disabled(), CommonPools.blockingTaskExecutor(),
SuccessFunction.always(), 0, multipartUploadsLocation, ImmutableList.of(),
HttpHeaders.of(), ctx -> RequestId.random(), serviceErrorHandler),
HttpHeaders.of(), ctx -> RequestId.random(), serviceErrorHandler,
NOOP_CONTEXT_HOOK),
new ServiceConfig(route2, route2,
SERVICE, defaultLogName, defaultServiceName, defaultServiceNaming, 0, 0,
false, AccessLogWriter.disabled(), CommonPools.blockingTaskExecutor(),
SuccessFunction.always(), 0, multipartUploadsLocation, ImmutableList.of(),
HttpHeaders.of(), ctx -> RequestId.random(), serviceErrorHandler));
HttpHeaders.of(), ctx -> RequestId.random(), serviceErrorHandler,
NOOP_CONTEXT_HOOK));
FALLBACK_SERVICE = new ServiceConfig(Route.ofCatchAll(), Route.ofCatchAll(), SERVICE,
defaultLogName, defaultServiceName,
defaultServiceNaming, 0, 0, false, AccessLogWriter.disabled(),
CommonPools.blockingTaskExecutor(),
SuccessFunction.always(), 0, multipartUploadsLocation,
ImmutableList.of(), HttpHeaders.of(), ctx -> RequestId.random(),
serviceErrorHandler);
serviceErrorHandler, NOOP_CONTEXT_HOOK);
HOST = new VirtualHost(
"localhost", "localhost", 0, null, SERVICES, FALLBACK_SERVICE, RejectedRouteHandler.DISABLED,
unused -> NOPLogger.NOP_LOGGER, defaultServiceNaming, defaultLogName, 0, 0, false,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,8 @@
package com.linecorp.armeria.client;

import static com.linecorp.armeria.client.ClientOptions.REDIRECT_CONFIG;
import static com.linecorp.armeria.internal.common.RequestContextUtil.NOOP_CONTEXT_HOOK;
import static com.linecorp.armeria.internal.common.RequestContextUtil.mergeHooks;
import static java.util.Objects.requireNonNull;

import java.time.Duration;
Expand All @@ -35,6 +37,8 @@
import com.linecorp.armeria.common.HttpRequest;
import com.linecorp.armeria.common.HttpResponse;
import com.linecorp.armeria.common.Request;
import com.linecorp.armeria.common.RequestContext;
import com.linecorp.armeria.common.RequestContextStorage;
import com.linecorp.armeria.common.RequestId;
import com.linecorp.armeria.common.SuccessFunction;
import com.linecorp.armeria.common.annotation.Nullable;
Expand All @@ -55,6 +59,7 @@ public class AbstractClientOptionsBuilder {

@Nullable
private Consumer<ClientRequestContext> contextCustomizer;
private Supplier<AutoCloseable> contextHook = NOOP_CONTEXT_HOOK;

/**
* Creates a new instance.
Expand Down Expand Up @@ -260,6 +265,19 @@ public AbstractClientOptionsBuilder endpointRemapper(
return option(ClientOptions.ENDPOINT_REMAPPER, endpointRemapper);
}

/**
* Sets the {@link Supplier} which provides an {@link AutoCloseable} and will be called whenever this
* {@link RequestContext} is popped from the {@link RequestContextStorage}.
*
* @param contextHook the {@link Supplier} that provides an {@link AutoCloseable}
*/
@UnstableApi
public AbstractClientOptionsBuilder contextHook(Supplier<? extends AutoCloseable> contextHook) {
minwoox marked this conversation as resolved.
Show resolved Hide resolved
minwoox marked this conversation as resolved.
Show resolved Hide resolved
requireNonNull(contextHook, "contextHook");
this.contextHook = mergeHooks(this.contextHook, contextHook);
return this;
}

/**
* Adds the specified HTTP-level {@code decorator}.
*
Expand Down Expand Up @@ -503,12 +521,13 @@ protected final ClientOptions buildOptions() {
protected final ClientOptions buildOptions(@Nullable ClientOptions baseOptions) {
final Collection<ClientOptionValue<?>> optVals = options.values();
final int numOpts = optVals.size();
final int extra = contextCustomizer == null ? 2 : 3;
final int extra = contextCustomizer == null ? 3 : 4;
final ClientOptionValue<?>[] optValArray = optVals.toArray(new ClientOptionValue[numOpts + extra]);
optValArray[numOpts] = ClientOptions.DECORATION.newValue(decoration.build());
optValArray[numOpts + 1] = ClientOptions.HEADERS.newValue(headers.build());
optValArray[numOpts + 2] = ClientOptions.CONTEXT_HOOK.newValue(contextHook);
if (contextCustomizer != null) {
optValArray[numOpts + 2] = ClientOptions.CONTEXT_CUSTOMIZER.newValue(contextCustomizer);
optValArray[numOpts + 3] = ClientOptions.CONTEXT_CUSTOMIZER.newValue(contextCustomizer);
}

if (baseOptions != null) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
package com.linecorp.armeria.client;

import static com.linecorp.armeria.internal.common.HttpHeadersUtil.CLOSE_STRING;
import static com.linecorp.armeria.internal.common.RequestContextUtil.NOOP_CONTEXT_HOOK;
import static java.util.Objects.requireNonNull;

import java.util.Arrays;
Expand All @@ -34,6 +35,8 @@
import com.linecorp.armeria.common.HttpHeaders;
import com.linecorp.armeria.common.HttpRequest;
import com.linecorp.armeria.common.HttpResponse;
import com.linecorp.armeria.common.RequestContext;
import com.linecorp.armeria.common.RequestContextStorage;
import com.linecorp.armeria.common.RequestId;
import com.linecorp.armeria.common.SessionProtocol;
import com.linecorp.armeria.common.SuccessFunction;
Expand Down Expand Up @@ -149,6 +152,10 @@ public final class ClientOptions
public static final ClientOption<Function<? super Endpoint, ? extends EndpointGroup>> ENDPOINT_REMAPPER =
ClientOption.define("ENDPOINT_REMAPPER", Function.identity());

@UnstableApi
public static final ClientOption<Supplier<? extends AutoCloseable>> CONTEXT_HOOK =
ikhoon marked this conversation as resolved.
Show resolved Hide resolved
ClientOption.define("CONTEXT_HOOK", NOOP_CONTEXT_HOOK);

private static final List<AsciiString> PROHIBITED_HEADER_NAMES = ImmutableList.of(
HttpHeaderNames.HTTP2_SETTINGS,
HttpHeaderNames.METHOD,
Expand Down Expand Up @@ -378,6 +385,16 @@ public Consumer<ClientRequestContext> contextCustomizer() {
return get(CONTEXT_CUSTOMIZER);
}

/**
* Returns the {@link Supplier} which provides an {@link AutoCloseable} and will be called whenever this
* {@link RequestContext} is popped from the {@link RequestContextStorage}.
*/
@UnstableApi
@SuppressWarnings("unchecked")
public Supplier<AutoCloseable> contextHook() {
return (Supplier<AutoCloseable>) get(CONTEXT_HOOK);
}

/**
* Returns a new {@link ClientOptionsBuilder} created from this {@link ClientOptions}.
*/
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
import java.time.Duration;
import java.util.concurrent.CompletableFuture;
import java.util.function.Consumer;
import java.util.function.Supplier;

import com.linecorp.armeria.client.endpoint.EndpointGroup;
import com.linecorp.armeria.common.ExchangeType;
Expand Down Expand Up @@ -156,6 +157,16 @@
return unwrap().exchangeType();
}

@Override
public void hook(Supplier<? extends AutoCloseable> contextHook) {
unwrap().hook(contextHook);
}

Check warning on line 163 in core/src/main/java/com/linecorp/armeria/client/ClientRequestContextWrapper.java

View check run for this annotation

Codecov / codecov/patch

core/src/main/java/com/linecorp/armeria/client/ClientRequestContextWrapper.java#L162-L163

Added lines #L162 - L163 were not covered by tests

@Override
public Supplier<AutoCloseable> hook() {
return unwrap().hook();

Check warning on line 167 in core/src/main/java/com/linecorp/armeria/client/ClientRequestContextWrapper.java

View check run for this annotation

Codecov / codecov/patch

core/src/main/java/com/linecorp/armeria/client/ClientRequestContextWrapper.java#L167

Added line #L167 was not covered by tests
}

@Override
public boolean isCancelled() {
return unwrap().isCancelled();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -172,6 +172,11 @@ public WebClientBuilder endpointRemapper(
return (WebClientBuilder) super.endpointRemapper(endpointRemapper);
}

@Override
public WebClientBuilder contextHook(Supplier<? extends AutoCloseable> contextHook) {
Copy link
Contributor

Choose a reason for hiding this comment

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

We probably need to override the rest of the methods

e.g.

RestClient.builder()
          .contextHook(() -> {})
          .build();

return (WebClientBuilder) super.contextHook(contextHook);
}

@Override
public WebClientBuilder decorator(
Function<? super HttpClient, ? extends HttpClient> decorator) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -538,6 +538,31 @@ default ByteBufAllocator alloc() {
@MustBeClosed
SafeCloseable push();

/**
* Adds a hook which is invoked whenever this {@link RequestContext} is pushed to the
* {@link RequestContextStorage}. The {@link AutoCloseable} returned by {@code contextHook} will be called
* whenever this {@link RequestContext} is popped from the {@link RequestContextStorage}.
* This method is useful when you need to propagate a custom context in this {@link RequestContext}'s scope.
*
* <p>Note:
* <ul>
* <li>The {@code contextHook} is not invoked when this {@link #hook(Supplier)} method is called
* thus you need to call it yourself if you want to apply the hook in the current thread. </li>
* <li>This operation is highly performance-sensitive operation, and thus it's not a good idea to run a
* time-consuming task.</li>
* </ul>
*/
@UnstableApi
void hook(Supplier<? extends AutoCloseable> contextHook);
minwoox marked this conversation as resolved.
Show resolved Hide resolved

/**
* Returns the hook which is invoked whenever this {@link RequestContext} is pushed to the
* {@link RequestContextStorage}. The {@link SafeCloseable} returned by the {@link Supplier} will be
* called whenever this {@link RequestContext} is popped from the {@link RequestContextStorage}.
*/
@UnstableApi
Supplier<AutoCloseable> hook();
minwoox marked this conversation as resolved.
Show resolved Hide resolved

@Override
default RequestContext unwrap() {
return (RequestContext) Unwrappable.super.unwrap();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -225,7 +225,7 @@ private DefaultClientRequestContext(
super(meterRegistry, desiredSessionProtocol(sessionProtocol, options), id, method, reqTarget,
firstNonNull(requestOptions.exchangeType(), ExchangeType.BIDI_STREAMING),
requestAutoAbortDelayMillis(options, requestOptions), req, rpcReq,
getAttributes(root));
getAttributes(root), options.contextHook());

this.eventLoop = eventLoop;
this.options = requireNonNull(options, "options");
Expand Down Expand Up @@ -500,7 +500,7 @@ private DefaultClientRequestContext(DefaultClientRequestContext ctx,
SessionProtocol sessionProtocol, HttpMethod method,
RequestTarget reqTarget) {
super(ctx.meterRegistry(), sessionProtocol, id, method, reqTarget, ctx.exchangeType(),
ctx.requestAutoAbortDelayMillis(), req, rpcReq, getAttributes(ctx.root()));
ctx.requestAutoAbortDelayMillis(), req, rpcReq, getAttributes(ctx.root()), ctx.hook());

// The new requests cannot be null if it was previously non-null.
if (ctx.request() != null) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,8 @@
package com.linecorp.armeria.internal.common;

import static com.google.common.base.MoreObjects.firstNonNull;
import static com.linecorp.armeria.internal.common.RequestContextUtil.NOOP_CONTEXT_HOOK;
import static com.linecorp.armeria.internal.common.RequestContextUtil.mergeHooks;
import static java.util.Objects.requireNonNull;

import java.util.Iterator;
Expand Down Expand Up @@ -72,7 +74,7 @@
private volatile HttpRequest req;
@Nullable
private volatile RpcRequest rpcReq;
@Nullable // Updated via `contextHookUpdater`
// Updated via `contextHookUpdater`
private volatile Supplier<AutoCloseable> contextHook;

/**
Expand All @@ -83,7 +85,7 @@
RequestId id, HttpMethod method, RequestTarget reqTarget, ExchangeType exchangeType,
long requestAutoAbortDelayMillis,
@Nullable HttpRequest req, @Nullable RpcRequest rpcReq,
@Nullable AttributesGetters rootAttributeMap) {
@Nullable AttributesGetters rootAttributeMap, Supplier<? extends AutoCloseable> contextHook) {
assert req != null || rpcReq != null;

this.meterRegistry = requireNonNull(meterRegistry, "meterRegistry");
Expand All @@ -102,6 +104,8 @@
originalRequest = firstNonNull(req, rpcReq);
this.req = req;
this.rpcReq = rpcReq;
//noinspection unchecked
this.contextHook = (Supplier<AutoCloseable>) contextHook;
}

@Override
Expand Down Expand Up @@ -267,30 +271,21 @@
@Override
public void hook(Supplier<? extends AutoCloseable> contextHook) {
requireNonNull(contextHook, "contextHook");

if (contextHook == NOOP_CONTEXT_HOOK) {
return;

Check warning on line 276 in core/src/main/java/com/linecorp/armeria/internal/common/NonWrappingRequestContext.java

View check run for this annotation

Codecov / codecov/patch

core/src/main/java/com/linecorp/armeria/internal/common/NonWrappingRequestContext.java#L276

Added line #L276 was not covered by tests
}

for (;;) {
final Supplier<? extends AutoCloseable> oldContextHook = this.contextHook;
final Supplier<? extends AutoCloseable> newContextHook;
if (oldContextHook == null) {
newContextHook = contextHook;
} else {
newContextHook = () -> {
final AutoCloseable oldHook = oldContextHook.get();
final AutoCloseable newHook = contextHook.get();
return () -> {
oldHook.close();
newHook.close();
};
};
}

final Supplier<? extends AutoCloseable> newContextHook = mergeHooks(oldContextHook, contextHook);
if (contextHookUpdater.compareAndSet(this, oldContextHook, newContextHook)) {
break;
}
}
}

@Override
@Nullable
public Supplier<AutoCloseable> hook() {
return contextHook;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -16,14 +16,9 @@

package com.linecorp.armeria.internal.common;

import java.util.function.Supplier;

import com.linecorp.armeria.common.AttributesGetters;
import com.linecorp.armeria.common.Request;
import com.linecorp.armeria.common.RequestContext;
import com.linecorp.armeria.common.RequestContextStorage;
import com.linecorp.armeria.common.annotation.Nullable;
import com.linecorp.armeria.common.util.SafeCloseable;

import io.netty.util.AttributeKey;

Expand All @@ -33,25 +28,6 @@
*/
public interface RequestContextExtension extends RequestContext {

/**
* Adds a hook which is invoked whenever this {@link NonWrappingRequestContext} is pushed to the
* {@link RequestContextStorage}. The {@link AutoCloseable} returned by {@code contextHook} will be called
* whenever this {@link RequestContext} is popped from the {@link RequestContextStorage}.
* This method is useful when you need to propagate a custom context in this {@link RequestContext}'s scope.
*
* <p>Note that this operation is highly performance-sensitive operation, and thus
* it's not a good idea to run a time-consuming task.
*/
void hook(Supplier<? extends AutoCloseable> contextHook);

/**
* Returns the hook which is invoked whenever this {@link NonWrappingRequestContext} is pushed to the
* {@link RequestContextStorage}. The {@link SafeCloseable} returned by the {@link Supplier} will be
* called whenever this {@link RequestContext} is popped from the {@link RequestContextStorage}.
*/
@Nullable
Supplier<AutoCloseable> hook();

/**
* Returns the {@link AttributesGetters} which stores the pairs of an {@link AttributeKey} and an object
* set via {@link #setAttr(AttributeKey, Object)}.
Expand Down
Loading