Add HTTP and gRPC engine ping support

README.md

@@ -68,6 +68,17 @@ ChalkClient client = ChalkClient.builder()
 ```
 
 
+## Ping
+
+Use `ping` to check connectivity to the Chalk engine. The engine echoes the
+number sent by the client. Calling `ping()` without an argument sends a random
+number between 0 and 999.
+
+```java
+int echoed = client.ping(3); // 3
+int randomEcho = client.ping();
+```
+
 ## Memory management
 Chalk uses Arrow objects during serialization of requests and deserialization of responses. These objects need to be
 closed to free up memory. 

src/main/java/ai/chalk/client/AuthenticatedHeaderClientInterceptor.java

@@ -47,7 +47,9 @@ public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall(
             public void start(Listener<RespT> responseListener, Metadata headers) {
                 GetTokenResponse token = tokenRefresher.getToken();
                 for (Map.Entry<Metadata.Key<String>, String> entry : allHeaders.entrySet()) {
-                    headers.put(entry.getKey(), entry.getValue());
+                    if (!headers.containsKey(entry.getKey())) {
+                        headers.put(entry.getKey(), entry.getValue());
+                    }
                 }
                 headers.put(GrpcHeaders.AUTHORIZATION_KEY, "Bearer " + token.getAccessToken());
                 super.start(responseListener, headers);

src/main/java/ai/chalk/client/ChalkClient.java

@@ -7,6 +7,7 @@
 import java.nio.file.Path;
 import java.time.Duration;
 import java.util.Optional;
+import java.util.concurrent.ThreadLocalRandom;
 
 public interface ChalkClient extends AutoCloseable {
 
@@ -62,6 +63,27 @@ static ChalkClient createGrpc() throws ChalkException {
         return builder().withGrpc().build();
     }
 
+    /**
+     * Pings the Chalk engine with a random number and returns the number echoed
+     * by the engine.
+     *
+     * @return the number echoed by the engine
+     * @throws ChalkException if the engine cannot be reached
+     */
+    default int ping() throws ChalkException {
+        return ping(ThreadLocalRandom.current().nextInt(1000));
+    }
+
+    /**
+     * Pings the Chalk engine with {@code num} and returns the number echoed by
+     * the engine.
+     *
+     * @param num the number to send to the engine
+     * @return the number echoed by the engine
+     * @throws ChalkException if the engine cannot be reached
+     */
+    int ping(int num) throws ChalkException;
+
     /**
      * OnlineQuery computes features values using online resolvers.
      * <p> See {@link OnlineQueryParams} for more details on the parameters.

src/main/java/ai/chalk/client/ChalkClientImpl.java

@@ -11,6 +11,8 @@
 import ai.chalk.internal.config.models.SourcedConfig;
 import ai.chalk.internal.request.RequestHandler;
 import ai.chalk.internal.request.models.OnlineQueryBulkResponse;
+import ai.chalk.internal.request.models.PingRequest;
+import ai.chalk.internal.request.models.PingResponse;
 import ai.chalk.internal.request.models.SendRequestParams;
 import ai.chalk.models.OnlineQueryParamsComplete;
 import ai.chalk.models.OnlineQueryResult;
@@ -74,6 +76,19 @@ public OnlineQueryResult onlineQuery(OnlineQueryParamsComplete params) throws Ch
         return onlineQueryFeather(params);
     }
 
+    @Override
+    public int ping(int num) throws ChalkException {
+        SendRequestParams request = new SendRequestParams.Builder(null)
+                .path("/ping")
+                .body(new PingRequest(num))
+                .method("POST")
+                .isEngineRequest(true)
+                .build();
+
+        HttpResponse<byte[]> response = this.handler.sendRequest(request);
+        return this.handler.deserializeResponseBody(response.body(), PingResponse.class).num();
+    }
+
     private OnlineQueryResult onlineQueryFeather(OnlineQueryParamsComplete params) throws ChalkException {
         byte[] bodyBytes;
         try (

src/main/java/ai/chalk/client/ConfigEnvVars.java

@@ -21,7 +21,12 @@ public class ConfigEnvVars {
      * This is the identifier that distinguishes between different
      * environments in your Chalk project.
      */
-    public static String environmentIdKey = "CHALK_ACTIVE_ENVIRONMENT";
+    public static String environmentIdKey = "CHALK_ENVIRONMENT";
+
+    /**
+     * Legacy environment variable supported for backwards compatibility.
+     */
+    public static String legacyEnvironmentIdKey = "CHALK_ACTIVE_ENVIRONMENT";
     /**
      * The path to the root CA certificate file.
      */

src/main/java/ai/chalk/client/GRPCClient.java

@@ -12,6 +12,7 @@
 import ai.chalk.models.UploadFeaturesParams;
 import ai.chalk.models.UploadFeaturesResult;
 import ai.chalk.protos.chalk.common.v1.*;
+import ai.chalk.protos.chalk.engine.v1.PingRequest;
 import ai.chalk.protos.chalk.engine.v1.QueryServiceGrpc;
 import ai.chalk.protos.chalk.server.v1.AuthServiceGrpc;
 import ai.chalk.protos.chalk.server.v1.GetTokenResponse;
@@ -122,6 +123,7 @@ public GRPCClient(BuilderImpl builder) throws ChalkException {
             }
         }
         engineHost = engineHost.replaceFirst("^https?://", "");
+        boolean queryThroughApiServer = engineHost.equals(grpcHost);
 
         /**
          * Create static retry policy (TODO: make it configurable)
@@ -145,8 +147,13 @@ public GRPCClient(BuilderImpl builder) throws ChalkException {
                 .maxInboundMessageSize(1024 * 1024 * 500)
                 .intercept(
                         new AuthenticatedHeaderClientInterceptor(
-                                ServerType.ENGINE,
-                                Map.of(),
+                                queryThroughApiServer ? ServerType.SERVER : ServerType.ENGINE,
+                                queryThroughApiServer
+                                        ? Map.of(
+                                                "x-chalk-env-id", resolvedEnvironmentId,
+                                                "x-chalk-deployment-type", "engine-grpc"
+                                        )
+                                        : Map.of("x-chalk-env-id", resolvedEnvironmentId),
                                 tokenRefresher,
                                 builder.getDeploymentTag()
                         )
@@ -188,6 +195,14 @@ private RequestHeaderInterceptor getRequestHeaderInterceptor(
         return new RequestHeaderInterceptor(environmentIdOverride, this.resolvedEnvironmentId, queryName);
     }
 
+    @Override
+    public int ping(int num) {
+        return this.stubsProvider.getQueryStub(Optional.empty())
+                .withInterceptors(this.getRequestHeaderInterceptor(null, null))
+                .ping(PingRequest.newBuilder().setNum(num).build())
+                .getNum();
+    }
+
 
     public OnlineQueryResult onlineQuery(OnlineQueryParamsComplete params) throws ChalkException {
         byte[] bodyBytes;

src/main/java/ai/chalk/client/ResolvedConfig.java

@@ -36,6 +36,7 @@ public static ResolvedConfig fromBuilder(
                 SourcedConfig.firstNonEmpty(
                         SourcedConfig.fromBuilder(builder.getEnvironmentId()),
                         SourcedConfig.fromEnvVar(ConfigEnvVars.environmentIdKey),
+                        SourcedConfig.fromEnvVar(ConfigEnvVars.legacyEnvironmentIdKey),
                         SourcedConfig.fromConfigFile(chalkYamlConfig.getActiveEnvironment())
                 ),
                 SourcedConfig.firstNonEmpty(

src/main/java/ai/chalk/internal/config/models/ProjectToken.java

@@ -1,12 +1,14 @@
 package ai.chalk.internal.config.models;
 
+import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
 import lombok.Data;
 import lombok.NoArgsConstructor;
 
 import java.time.LocalDateTime;
 
 @Data
 @NoArgsConstructor
+@JsonIgnoreProperties(ignoreUnknown = true)
 public class ProjectToken {
     private String name;
     private String clientId;

src/main/java/ai/chalk/internal/request/models/PingRequest.java

@@ -0,0 +1,4 @@
+package ai.chalk.internal.request.models;
+
+public record PingRequest(int num) {
+}

src/main/java/ai/chalk/internal/request/models/PingResponse.java

@@ -0,0 +1,4 @@
+package ai.chalk.internal.request.models;
+
+public record PingResponse(int num) {
+}

src/test/java/ai/chalk/client/TestPing.java

@@ -0,0 +1,154 @@
+package ai.chalk.client;
+
+import ai.chalk.protos.chalk.engine.v1.PingRequest;
+import ai.chalk.protos.chalk.engine.v1.PingResponse;
+import ai.chalk.protos.chalk.engine.v1.QueryServiceGrpc;
+import ai.chalk.protos.chalk.server.v1.AuthServiceGrpc;
+import com.google.protobuf.Timestamp;
+import com.sun.net.httpserver.HttpExchange;
+import com.sun.net.httpserver.HttpServer;
+import io.grpc.Server;
+import io.grpc.Metadata;
+import io.grpc.ServerInterceptors;
+import io.grpc.ServerCall;
+import io.grpc.ServerCallHandler;
+import io.grpc.ServerInterceptor;
+import io.grpc.netty.shaded.io.grpc.netty.NettyServerBuilder;
+import io.grpc.stub.StreamObserver;
+import org.junit.jupiter.api.Test;
+
+import java.io.IOException;
+import java.net.InetSocketAddress;
+import java.nio.charset.StandardCharsets;
+import java.time.Instant;
+import java.util.concurrent.atomic.AtomicReference;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+public class TestPing {
+    private static final String ENVIRONMENT_ID = "test-environment";
+
+    @Test
+    public void testHttpPing() throws Exception {
+        AtomicReference<String> pingBody = new AtomicReference<>();
+        HttpServer server = HttpServer.create(new InetSocketAddress("localhost", 0), 0);
+        server.createContext("/v1/oauth/token", exchange -> respondJson(
+                exchange,
+                """
+                {
+                  "access_token": "test-token",
+                  "expires_in": 3600,
+                  "primary_environment": "%s",
+                  "engines": {}
+                }
+                """.formatted(ENVIRONMENT_ID)
+        ));
+        server.createContext("/ping", exchange -> {
+            pingBody.set(new String(exchange.getRequestBody().readAllBytes(), StandardCharsets.UTF_8));
+            respondJson(exchange, "{\"num\":37}");
+        });
+        server.start();
+
+        String serverUrl = "http://localhost:" + server.getAddress().getPort();
+        try (ChalkClient client = ChalkClient.builder()
+                .withApiServer(serverUrl)
+                .withQueryServer(serverUrl)
+                .withClientId("test-client")
+                .withClientSecret("test-secret")
+                .withEnvironmentId(ENVIRONMENT_ID)
+                .build()) {
+            assertEquals(37, client.ping(37));
+            assertEquals("{\"num\":37}", pingBody.get());
+        } finally {
+            server.stop(0);
+        }
+    }
+
+    @Test
+    public void testGrpcPing() throws Exception {
+        AtomicReference<String> environmentHeader = new AtomicReference<>();
+        AtomicReference<String> serverHeader = new AtomicReference<>();
+        AtomicReference<String> deploymentTypeHeader = new AtomicReference<>();
+        Metadata.Key<String> environmentHeaderKey = Metadata.Key.of(
+                "x-chalk-env-id",
+                Metadata.ASCII_STRING_MARSHALLER
+        );
+        Metadata.Key<String> serverHeaderKey = Metadata.Key.of(
+                "x-chalk-server",
+                Metadata.ASCII_STRING_MARSHALLER
+        );
+        Metadata.Key<String> deploymentTypeHeaderKey = Metadata.Key.of(
+                "x-chalk-deployment-type",
+                Metadata.ASCII_STRING_MARSHALLER
+        );
+        Server server = NettyServerBuilder.forPort(0)
+                .addService(new AuthServiceGrpc.AuthServiceImplBase() {
+                    @Override
+                    public void getToken(
+                            ai.chalk.protos.chalk.server.v1.GetTokenRequest request,
+                            StreamObserver<ai.chalk.protos.chalk.server.v1.GetTokenResponse> responseObserver
+                    ) {
+                        responseObserver.onNext(ai.chalk.protos.chalk.server.v1.GetTokenResponse.newBuilder()
+                                .setAccessToken("test-token")
+                                .setExpiresAt(Timestamp.newBuilder()
+                                        .setSeconds(Instant.now().plusSeconds(3600).getEpochSecond()))
+                                .setPrimaryEnvironment(ENVIRONMENT_ID)
+                                .putEnvironmentIdToName(ENVIRONMENT_ID, "test")
+                                .build());
+                        responseObserver.onCompleted();
+                    }
+                })
+                .addService(ServerInterceptors.intercept(
+                        new QueryServiceGrpc.QueryServiceImplBase() {
+                            @Override
+                            public void ping(PingRequest request, StreamObserver<PingResponse> responseObserver) {
+                                responseObserver.onNext(PingResponse.newBuilder().setNum(request.getNum()).build());
+                                responseObserver.onCompleted();
+                            }
+                        },
+                        new ServerInterceptor() {
+                            @Override
+                            public <ReqT, RespT> ServerCall.Listener<ReqT> interceptCall(
+                                    ServerCall<ReqT, RespT> call,
+                                    Metadata headers,
+                                    ServerCallHandler<ReqT, RespT> next
+                            ) {
+                                environmentHeader.set(headers.get(environmentHeaderKey));
+                                serverHeader.set(headers.get(serverHeaderKey));
+                                deploymentTypeHeader.set(headers.get(deploymentTypeHeaderKey));
+                                return next.startCall(call, headers);
+                            }
+                        }
+                ))
+                .build()
+                .start();
+
+        String grpcHost = "localhost:" + server.getPort();
+        try (ChalkClient client = ChalkClient.builder()
+                .withApiServer("http://" + grpcHost)
+                .withQueryServer(grpcHost)
+                .withClientId("test-client")
+                .withClientSecret("test-secret")
+                .withEnvironmentId(ENVIRONMENT_ID)
+                .withGrpc()
+                .build()) {
+            assertEquals(41, client.ping(41));
+            assertEquals(ENVIRONMENT_ID, environmentHeader.get());
+            assertEquals("go-api", serverHeader.get());
+            assertEquals("engine-grpc", deploymentTypeHeader.get());
+            int randomEcho = client.ping();
+            assertTrue(randomEcho >= 0 && randomEcho < 1000);
+        } finally {
+            server.shutdownNow();
+        }
+    }
+
+    private static void respondJson(HttpExchange exchange, String body) throws IOException {
+        byte[] response = body.getBytes(StandardCharsets.UTF_8);
+        exchange.getResponseHeaders().set("Content-Type", "application/json");
+        exchange.sendResponseHeaders(200, response.length);
+        exchange.getResponseBody().write(response);
+        exchange.close();
+    }
+}

src/test/java/ai/chalk/internal/TestConfigLoader.java

@@ -23,6 +23,7 @@ public Path writeTempYAMLFile() throws Exception {
                 validUntil: 2025-05-02T17:32:46.165000
                 apiServer: https://api.chalk.ai
                 activeEnvironment: myprojectenvid
+                activeEnvironmentName: My Project
               default:
                 name: Default Token
                 clientId: client-default
@@ -54,4 +55,3 @@ public void TestLoadAllTokens() throws Exception {
         assert defaultToken.getValidUntil().equals(LocalDateTime.of(2024, 9, 11, 15, 7, 50, 812_000_000));
     }
 }
-