From 383c0d6ae0ea334940ec67d7e7b70a8d1aae25da Mon Sep 17 00:00:00 2001 From: Sahej Singh Date: Thu, 27 Aug 2026 15:36:45 -0400 Subject: [PATCH] Add engine ping support --- README.md | 11 ++ .../AuthenticatedHeaderClientInterceptor.java | 4 +- .../java/ai/chalk/client/ChalkClient.java | 22 +++ .../java/ai/chalk/client/ChalkClientImpl.java | 15 ++ .../java/ai/chalk/client/ConfigEnvVars.java | 7 +- src/main/java/ai/chalk/client/GRPCClient.java | 19 ++- .../java/ai/chalk/client/ResolvedConfig.java | 1 + .../internal/config/models/ProjectToken.java | 2 + .../internal/request/models/PingRequest.java | 4 + .../internal/request/models/PingResponse.java | 4 + src/test/java/ai/chalk/client/TestPing.java | 154 ++++++++++++++++++ .../ai/chalk/internal/TestConfigLoader.java | 2 +- 12 files changed, 240 insertions(+), 5 deletions(-) create mode 100644 src/main/java/ai/chalk/internal/request/models/PingRequest.java create mode 100644 src/main/java/ai/chalk/internal/request/models/PingResponse.java create mode 100644 src/test/java/ai/chalk/client/TestPing.java diff --git a/README.md b/README.md index 6184b962..37816972 100644 --- a/README.md +++ b/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. diff --git a/src/main/java/ai/chalk/client/AuthenticatedHeaderClientInterceptor.java b/src/main/java/ai/chalk/client/AuthenticatedHeaderClientInterceptor.java index 5008b760..e610bf88 100644 --- a/src/main/java/ai/chalk/client/AuthenticatedHeaderClientInterceptor.java +++ b/src/main/java/ai/chalk/client/AuthenticatedHeaderClientInterceptor.java @@ -47,7 +47,9 @@ public ClientCall interceptCall( public void start(Listener responseListener, Metadata headers) { GetTokenResponse token = tokenRefresher.getToken(); for (Map.Entry, 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); diff --git a/src/main/java/ai/chalk/client/ChalkClient.java b/src/main/java/ai/chalk/client/ChalkClient.java index 5a024b84..3ba53d32 100644 --- a/src/main/java/ai/chalk/client/ChalkClient.java +++ b/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. *

See {@link OnlineQueryParams} for more details on the parameters. diff --git a/src/main/java/ai/chalk/client/ChalkClientImpl.java b/src/main/java/ai/chalk/client/ChalkClientImpl.java index 96c7bf96..1097778b 100644 --- a/src/main/java/ai/chalk/client/ChalkClientImpl.java +++ b/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 response = this.handler.sendRequest(request); + return this.handler.deserializeResponseBody(response.body(), PingResponse.class).num(); + } + private OnlineQueryResult onlineQueryFeather(OnlineQueryParamsComplete params) throws ChalkException { byte[] bodyBytes; try ( diff --git a/src/main/java/ai/chalk/client/ConfigEnvVars.java b/src/main/java/ai/chalk/client/ConfigEnvVars.java index e7334173..dc91cde6 100644 --- a/src/main/java/ai/chalk/client/ConfigEnvVars.java +++ b/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. */ diff --git a/src/main/java/ai/chalk/client/GRPCClient.java b/src/main/java/ai/chalk/client/GRPCClient.java index 09aaac07..f2cdeaae 100644 --- a/src/main/java/ai/chalk/client/GRPCClient.java +++ b/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; diff --git a/src/main/java/ai/chalk/client/ResolvedConfig.java b/src/main/java/ai/chalk/client/ResolvedConfig.java index bd891f1c..8480cb5d 100644 --- a/src/main/java/ai/chalk/client/ResolvedConfig.java +++ b/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( diff --git a/src/main/java/ai/chalk/internal/config/models/ProjectToken.java b/src/main/java/ai/chalk/internal/config/models/ProjectToken.java index f00b2871..fa0fe89d 100644 --- a/src/main/java/ai/chalk/internal/config/models/ProjectToken.java +++ b/src/main/java/ai/chalk/internal/config/models/ProjectToken.java @@ -1,5 +1,6 @@ package ai.chalk.internal.config.models; +import com.fasterxml.jackson.annotation.JsonIgnoreProperties; import lombok.Data; import lombok.NoArgsConstructor; @@ -7,6 +8,7 @@ @Data @NoArgsConstructor +@JsonIgnoreProperties(ignoreUnknown = true) public class ProjectToken { private String name; private String clientId; diff --git a/src/main/java/ai/chalk/internal/request/models/PingRequest.java b/src/main/java/ai/chalk/internal/request/models/PingRequest.java new file mode 100644 index 00000000..1053c2d7 --- /dev/null +++ b/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) { +} diff --git a/src/main/java/ai/chalk/internal/request/models/PingResponse.java b/src/main/java/ai/chalk/internal/request/models/PingResponse.java new file mode 100644 index 00000000..98dc6350 --- /dev/null +++ b/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) { +} diff --git a/src/test/java/ai/chalk/client/TestPing.java b/src/test/java/ai/chalk/client/TestPing.java new file mode 100644 index 00000000..2f9288ed --- /dev/null +++ b/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 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 environmentHeader = new AtomicReference<>(); + AtomicReference serverHeader = new AtomicReference<>(); + AtomicReference deploymentTypeHeader = new AtomicReference<>(); + Metadata.Key environmentHeaderKey = Metadata.Key.of( + "x-chalk-env-id", + Metadata.ASCII_STRING_MARSHALLER + ); + Metadata.Key serverHeaderKey = Metadata.Key.of( + "x-chalk-server", + Metadata.ASCII_STRING_MARSHALLER + ); + Metadata.Key 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 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 responseObserver) { + responseObserver.onNext(PingResponse.newBuilder().setNum(request.getNum()).build()); + responseObserver.onCompleted(); + } + }, + new ServerInterceptor() { + @Override + public ServerCall.Listener interceptCall( + ServerCall call, + Metadata headers, + ServerCallHandler 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(); + } +} diff --git a/src/test/java/ai/chalk/internal/TestConfigLoader.java b/src/test/java/ai/chalk/internal/TestConfigLoader.java index d8bf198f..bcf6c8c6 100644 --- a/src/test/java/ai/chalk/internal/TestConfigLoader.java +++ b/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)); } } -