Context propagation - day2

import java.util.concurrent.CompletableFuture;

import com.linecorp.armeria.client.HttpClient;
import com.linecorp.armeria.client.HttpClientBuilder;
import com.linecorp.armeria.common.AggregatedHttpResponse;
import com.linecorp.armeria.common.HttpRequest;
import com.linecorp.armeria.common.HttpResponse;
import com.linecorp.armeria.server.HttpService;
import com.linecorp.armeria.server.Server;
import com.linecorp.armeria.server.ServiceRequestContext;

import io.netty.util.AttributeKey;

public class RequestContextPropagation {

    private static final AttributeKey<String> MY_ID =
            AttributeKey.valueOf(RequestContextPropagation.class, "myId");

    public static void main(String[] args) {
        final Server backend1 = Server.builder()
                                      .http(8081)
                                      .route()
                                      .get("/api")
                                      .build((ctx, req) -> HttpResponse.of(200))
                                      .build();
        backend1.start().join();

        final Server backend2 = Server.builder()
                                      .http(8082)
                                      .route()
                                      .get("/api")
                                      .build((ctx, req) -> HttpResponse.of(200))
                                      .build();
        backend2.start().join();

        final HttpClient client1 = new HttpClientBuilder("h2c://127.0.0.1:8081")
                .decorator((delegate, ctx, req) -> {
                    System.err.println(ctx.hasAttr(MY_ID));
                    return delegate.execute(ctx, req);
                }).build();

        final HttpClient client2 = new HttpClientBuilder("h2c://127.0.0.1:8082")
                .decorator((delegate, ctx, req) -> {
                    System.err.println(ctx.hasAttr(MY_ID));
                    return delegate.execute(ctx, req);
                }).build();

        final Server server =
                Server.builder()
                      .http(8080)
                      .decorator((delegate, ctx, req) -> {
                          ctx.onChild((src, dst) -> {
                              dst.attr(MY_ID).set(src.attr(MY_ID).get());
                          });
                          ctx.attr(MY_ID).set("hello!!");
                          return delegate.serve(ctx, req);
                      })
                      .service("/service", new HttpService() {
                          @Override
                          public HttpResponse serve(ServiceRequestContext ctx, HttpRequest req)
                                  throws Exception {
                              final CompletableFuture<HttpResponse> future = new CompletableFuture<>();
                              final HttpResponse res = HttpResponse.from(future);

                              final CompletableFuture<AggregatedHttpResponse> client1Future =
                                      client1.get("/api").aggregate();

                              ctx.makeContextAware(client1Future).whenComplete((aggregatedRes1, cause1) -> {
                                  client2.get("/api").aggregate().thenApply(aggregatedRes2 -> {
                                      future.complete(HttpResponse.of(200));
                                      return null;
                                  });
                              });
                              return res;
                          }
                      }).build();
        server.start().join();

        sendRequest();
    }

    private static void sendRequest() {
        final HttpClient client = HttpClient.of("http://127.0.0.1:8080");
        final HttpResponse httpResponse = client.get("/service");
        final AggregatedHttpResponse res = httpResponse.aggregate().join();
        System.err.println(res.headers());
    }
}