Migrate from okhttp to armeria

benchmarks/src/main/java/zipkin2/elasticsearch/internal/BulkRequestBenchmarks.java

@@ -61,15 +61,15 @@ public class BulkRequestBenchmarks {
   @Benchmark public void buildAndWriteRequest_singleSpan() throws IOException {
     BulkCallBuilder builder = new BulkCallBuilder(es, 6.7f, "index-span");
     builder.index(spanIndex, "span", CLIENT_SPAN, BulkIndexWriter.SPAN);
-    builder.build().call.request().body().writeTo(Okio.buffer(Okio.blackhole()));
+    Okio.buffer(Okio.blackhole()).write(builder.build().request.content().array());
   }
 
   @Benchmark public void buildAndWriteRequest_tenSpans() throws IOException {
     BulkCallBuilder builder = new BulkCallBuilder(es, 6.7f, "index-span");
     for (int i = 0; i < 10; i++) {
       builder.index(spanIndex, "span", CLIENT_SPAN, BulkIndexWriter.SPAN);
     }
-    builder.build().call.request().body().writeTo(Okio.buffer(Okio.blackhole()));
+    Okio.buffer(Okio.blackhole()).write(builder.build().request.content().array());
   }
 
   // Convenience main entry-point

zipkin-server/pom.xml

@@ -190,17 +190,6 @@
       <version>${project.version}</version>
       <optional>true</optional>
     </dependency>
-    <dependency>
-      <groupId>com.squareup.okhttp3</groupId>
-      <artifactId>okhttp</artifactId>
-      <optional>true</optional>
-    </dependency>
-    <dependency>
-      <groupId>com.squareup.okhttp3</groupId>
-      <artifactId>logging-interceptor</artifactId>
-      <version>${okhttp.version}</version>
-      <optional>true</optional>
-    </dependency>
 
     <!-- MySQL backend -->
     <dependency>
@@ -295,6 +284,12 @@
       <optional>true</optional>
     </dependency>
 
+    <dependency>
+      <groupId>com.linecorp.armeria</groupId>
+      <artifactId>armeria-testing-junit4</artifactId>
+      <scope>test</scope>
+    </dependency>
+
     <dependency>
       <groupId>org.springframework.boot</groupId>
       <artifactId>spring-boot-starter-test</artifactId>
@@ -309,11 +304,6 @@
       <scope>test</scope>
     </dependency>
 
-    <dependency>
-      <groupId>com.squareup.okhttp3</groupId>
-      <artifactId>mockwebserver</artifactId>
-      <scope>test</scope>
-    </dependency>
     <dependency>
       <groupId>com.jayway.jsonpath</groupId>
       <artifactId>json-path</artifactId>

zipkin-server/src/main/java/zipkin2/server/internal/MetricsHealthController.java

@@ -17,9 +17,9 @@
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.fasterxml.jackson.databind.node.JsonNodeFactory;
 import com.fasterxml.jackson.databind.node.ObjectNode;
+import com.linecorp.armeria.common.AggregatedHttpResponse;
 import com.linecorp.armeria.common.HttpData;
 import com.linecorp.armeria.common.HttpHeaderNames;
-import com.linecorp.armeria.common.HttpResponse;
 import com.linecorp.armeria.common.MediaType;
 import com.linecorp.armeria.common.ResponseHeaders;
 import com.linecorp.armeria.server.annotation.Get;
@@ -85,7 +85,7 @@ public ObjectNode fetchMetricsFromMicrometer() {
   // Delegates the health endpoint from the Actuator to the root context path and can be deprecated
   // in future in favour of Actuator endpoints
   @Get("/health")
-  public HttpResponse getHealth() throws JsonProcessingException {
+  public AggregatedHttpResponse getHealth() throws JsonProcessingException {
     Health health = healthEndpoint.health();
 
     Map<String, Object> healthJson = new LinkedHashMap<>();
@@ -96,6 +96,6 @@ public HttpResponse getHealth() throws JsonProcessingException {
     ResponseHeaders headers = ResponseHeaders.builder(statusMapper.mapStatus(health.getStatus()))
       .contentType(MediaType.JSON)
       .setInt(HttpHeaderNames.CONTENT_LENGTH, body.length).build();
-    return HttpResponse.of(headers, HttpData.wrap(body));
+    return AggregatedHttpResponse.of(headers, HttpData.wrap(body));
   }
 }

zipkin-server/src/main/java/zipkin2/server/internal/elasticsearch/BasicAuthInterceptor.java

@@ -13,47 +13,73 @@
  */
 package zipkin2.server.internal.elasticsearch;
 
+import com.linecorp.armeria.client.Client;
+import com.linecorp.armeria.client.ClientRequestContext;
+import com.linecorp.armeria.client.SimpleDecoratingClient;
+import com.linecorp.armeria.common.HttpData;
+import com.linecorp.armeria.common.HttpHeaderNames;
+import com.linecorp.armeria.common.HttpRequest;
+import com.linecorp.armeria.common.HttpResponse;
+import com.linecorp.armeria.common.HttpStatus;
+import com.linecorp.armeria.common.util.Exceptions;
 import com.squareup.moshi.JsonReader;
+import io.netty.buffer.ByteBufHolder;
+import io.netty.util.ReferenceCountUtil;
 import java.io.IOException;
-import okhttp3.Credentials;
-import okhttp3.Interceptor;
-import okhttp3.Request;
-import okhttp3.Response;
-import okhttp3.ResponseBody;
+import java.io.UncheckedIOException;
+import java.nio.ByteBuffer;
+import java.nio.charset.StandardCharsets;
+import java.util.Base64;
+import okio.Okio;
+import zipkin2.internal.ReadBuffer;
 
 import static zipkin2.elasticsearch.internal.JsonReaders.enterPath;
 
 /**
  * Adds basic auth username and password to every request per
  * https://www.elastic.co/guide/en/x-pack/current/how-security-works.html
  */
-final class BasicAuthInterceptor implements Interceptor {
+final class BasicAuthInterceptor extends SimpleDecoratingClient<HttpRequest, HttpResponse> {
 
   private String basicCredentials;
 
-  BasicAuthInterceptor(ZipkinElasticsearchStorageProperties es) {
-    basicCredentials = Credentials.basic(es.getUsername(), es.getPassword());
+  BasicAuthInterceptor(
+    Client<HttpRequest, HttpResponse> client,
+    ZipkinElasticsearchStorageProperties es) {
+    super(client);
+    String token = es.getUsername() + ':' + es.getPassword();
+    basicCredentials = Base64.getEncoder().encodeToString(token.getBytes(StandardCharsets.UTF_8));
   }
 
-  @Override
-  public Response intercept(Chain chain) throws IOException {
-
-    Request input = chain.request();
-    Request requestWithCredentials = appendBasicAuthHeaderParameters(input);
-    Response response = chain.proceed(requestWithCredentials);
-    if (response.code() == 403) {
-      try (ResponseBody body = response.body()) {
-        JsonReader message = enterPath(JsonReader.of(body.source()), "message");
-        if (message != null) throw new IllegalStateException(message.nextString());
-      }
-      throw new IllegalStateException(response.toString());
-    }
-    return response;
-  }
-
-  private Request appendBasicAuthHeaderParameters(Request input) {
-
-    Request.Builder builder = input.newBuilder();
-    return builder.header("authorization", basicCredentials).build();
+  @Override public HttpResponse execute(ClientRequestContext ctx, HttpRequest req)
+    throws Exception {
+    ctx.addAdditionalRequestHeader(HttpHeaderNames.AUTHORIZATION, basicCredentials);
+    return HttpResponse.from(
+      delegate().execute(ctx, req).aggregateWithPooledObjects(ctx.eventLoop(), ctx.alloc())
+        .thenApply(msg -> {
+          HttpData content = msg.content();
+          if (!msg.status().equals(HttpStatus.FORBIDDEN) || content.isEmpty()) {
+            return HttpResponse.of(msg);
+          }
+          try {
+            final ByteBuffer buf;
+            if (content instanceof ByteBufHolder) {
+              buf = ((ByteBufHolder) content).content().nioBuffer();
+            } else {
+              buf = ByteBuffer.wrap(content.array());
+            }
+            try {
+              JsonReader message = enterPath(JsonReader.of(
+                Okio.buffer(Okio.source(ReadBuffer.wrapUnsafe(buf)))), "message");
+              if (message != null) throw new IllegalStateException(message.nextString());
+            } catch (IOException e) {
+              Exceptions.throwUnsafely(e);
+              throw new UncheckedIOException(e);  // unreachable
+            }
+            throw new IllegalStateException(msg.toString());
+          } finally {
+            ReferenceCountUtil.safeRelease(content);
+          }
+        }));
   }
 }