Skip to content

Commit 0450ccb

Browse files
sjmignotclaude
andcommitted
Address review: deadline, environment and version resolution
Fixes found reviewing the multi-query settings resolvers. - resolveDeadline dropped the client-level timeout as soon as any one query set its own. A query that sets no timeout is bounded by the client-level timeout, so that is its effective timeout and it has to take part in the maximum; otherwise a client built with a 5m timeout batching a 1s query alongside an unset one deadlined the whole request at 1s, killing a query that would have had 5m on its own. This contradicted the method's own stated invariant. A query that is effectively unbounded now yields no deadline at all. - resolveEnvironment compared raw per-query values, so naming the client's own environment on one query and leaving it unset on another was reported as a conflict even though both resolve to the same environment. It now compares the effective value, as resolveBranch already did. - Dropped the queryNameVersion agreement requirement. The justification does not hold for the version: there is no x-chalk-query-name-version header, and each sub-request carries its own version in its OnlineQueryContext, so queries differing only by version are satisfiable. - unmarshal() NPE'd on a null scalars table, which is exactly what a sub-query that failed to execute leaves behind, so the documented empty slot made the README's own pattern fail with an opaque error. It now explains where to find the reason. - Guard the varargs overload against a null array so it reports the same ClientException as the List overload rather than a raw NPE, and drop a repeated format argument that made the count-mismatch message tautological. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
1 parent 4375ce9 commit 0450ccb

6 files changed

Lines changed: 183 additions & 66 deletions

File tree

‎README.md‎

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -418,6 +418,10 @@ try (OnlineQueryMultiResult multi = client.onlineQueryMulti(List.of(users, merch
418418
// this query failed; the others are unaffected
419419
continue;
420420
}
421+
if (result.getScalarsTable() == null) {
422+
// this query never ran -- see getGlobalErrors() above for the reason
423+
continue;
424+
}
421425
// do something with the result
422426
}
423427
}
@@ -437,8 +441,10 @@ A few things to know:
437441
- **Close the wrapper, not the results.** Closing the `OnlineQueryMultiResult` closes every one of
438442
its results and releases all of their Arrow memory.
439443
- **Some settings apply to the whole request,** which carries one deadline and one set of headers.
440-
Every query must agree on `environmentId`, `branch`, `queryName` and `queryNameVersion`, or the
441-
call throws `ClientException` — a request resolves one branch to one deployment and routes named
442-
queries by an exact match on the query name, so a mixed batch could not be honored. Run those as
443-
separate queries. The deadline is the longest of the per-query timeouts, falling back to the
444+
Every query must agree on `environmentId`, `branch` and `queryName`, or the call throws
445+
`ClientException` — a request resolves one branch to one deployment and routes named queries by
446+
an exact match on the query name, so a mixed batch could not be honored. Run those as separate
447+
queries. Leaving one of these unset just takes the client's value, which is not a disagreement,
448+
and `queryNameVersion` may differ freely since each query carries its own. The deadline is the
449+
longest of the effective per-query timeouts, where a query that sets none is bounded by the
444450
client-level timeout, so per-query timeouts are not independently enforced.

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

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -188,15 +188,17 @@ default int ping() throws ChalkException {
188188
* <p> <b>Some settings apply to the whole request.</b> A request carries one
189189
* deadline and one set of headers, so:
190190
* <ul>
191-
* <li>Every query must agree on {@code environmentId}, {@code branch},
192-
* {@code queryName} and {@code queryNameVersion}. A disagreement throws
191+
* <li>Every query must agree on {@code environmentId}, {@code branch} and
192+
* {@code queryName}. A disagreement throws
193193
* {@link ai.chalk.exceptions.ClientException}, because a request resolves
194194
* one branch to one deployment and routes named queries by an exact match
195195
* on the query name, so a mixed batch could not be honored. Run those as
196-
* separate queries.</li>
197-
* <li>The deadline is the longest of the per-query timeouts, falling back to
198-
* the client-level timeout. Per-query timeouts are not independently
199-
* enforced.</li>
196+
* separate queries. A query that leaves one of these unset simply takes the
197+
* client's value, which is not a disagreement. {@code queryNameVersion} may
198+
* differ freely: each query carries its own.</li>
199+
* <li>The deadline is the longest of the effective per-query timeouts, where a
200+
* query that sets none is bounded by the client-level timeout. Per-query
201+
* timeouts are not independently enforced.</li>
200202
* <li>Branch selection is currently inert on the gRPC path for both single and
201203
* multi queries: the client does not yet send the header the API server
202204
* routes branches on.</li>
@@ -228,7 +230,7 @@ default OnlineQueryMultiResult onlineQueryMulti(List<OnlineQueryParamsComplete>
228230
*/
229231
default OnlineQueryMultiResult onlineQueryMulti(OnlineQueryParamsComplete... params)
230232
throws ChalkException {
231-
return onlineQueryMulti(Arrays.asList(params));
233+
return onlineQueryMulti(params == null ? null : Arrays.asList(params));
232234
}
233235

234236
/**

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

Lines changed: 40 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -240,8 +240,8 @@ public OnlineQueryMultiResult onlineQueryMulti(List<OnlineQueryParamsComplete> p
240240

241241
// Resolve the request-level settings before serializing anything, so a batch that
242242
// cannot be satisfied fails before doing the work of encoding its inputs.
243-
Duration deadline = resolveDeadline(paramsList).orElse(null);
244-
String environmentId = resolveEnvironment(paramsList);
243+
Duration deadline = resolveDeadline(paramsList, this.timeout).orElse(null);
244+
String environmentId = resolveEnvironment(paramsList, this.resolvedEnvironmentId);
245245
String queryName = resolveQueryName(paramsList);
246246
// Validated for agreement but deliberately unused: the gRPC client does not send the
247247
// "x-chalk-branch-id" header that the API server routes branches on, so branch
@@ -514,8 +514,7 @@ static OnlineQueryMultiResult decodeMultiResponse(
514514
// mismatch leaves every result unattributable. Fail rather than guess.
515515
if (response.getResponsesCount() != expectedCount) {
516516
throw new ClientException(String.format(
517-
"Expected %d sub-responses for %d queries, but got %d%s",
518-
expectedCount,
517+
"Expected %d sub-responses, one per query, but got %d%s",
519518
expectedCount,
520519
response.getResponsesCount(),
521520
globalErrors.length == 0 ? "" : "; errors: " + describeErrors(globalErrors)
@@ -587,18 +586,30 @@ static void validateMultiParams(List<OnlineQueryParamsComplete> paramsList) thro
587586
}
588587

589588
/**
590-
* Resolves the deadline for the whole request as the longest of the per-query
591-
* timeouts, so that no sub-query is cancelled before the time its caller
592-
* allowed. Returns empty when no query sets one, which lets the client-level
593-
* timeout apply. Per-query timeouts are not independently enforced: a single
589+
* Resolves the deadline for the whole request as the longest of the effective
590+
* per-query timeouts, so that no sub-query is cancelled before the time its
591+
* caller allowed. Per-query timeouts are not independently enforced: a single
594592
* request carries a single deadline.
593+
*
594+
* <p> A query that sets no timeout of its own is bounded by the client-level
595+
* timeout, so that is its effective timeout and it takes part in the maximum.
596+
* Leaving it out would let one query's short timeout cut short a query that was
597+
* relying on a longer client-level default. If any query is effectively
598+
* unbounded -- it sets no timeout and the client sets none either -- the request
599+
* gets no deadline.
595600
*/
596-
static Optional<Duration> resolveDeadline(List<OnlineQueryParamsComplete> paramsList) {
601+
static Optional<Duration> resolveDeadline(
602+
List<OnlineQueryParamsComplete> paramsList,
603+
Optional<Duration> clientLevelTimeout
604+
) {
597605
Duration longest = null;
598606
for (OnlineQueryParamsComplete params : paramsList) {
599-
Duration timeout = params.getTimeout();
607+
Duration timeout = params.getTimeout() != null
608+
? params.getTimeout()
609+
: clientLevelTimeout.orElse(null);
600610
if (timeout == null) {
601-
continue;
611+
// This query is unbounded, so the request cannot carry a deadline.
612+
return Optional.empty();
602613
}
603614
if (longest == null || timeout.compareTo(longest) > 0) {
604615
longest = timeout;
@@ -607,9 +618,17 @@ static Optional<Duration> resolveDeadline(List<OnlineQueryParamsComplete> params
607618
return Optional.ofNullable(longest);
608619
}
609620

610-
static @Nullable String resolveEnvironment(List<OnlineQueryParamsComplete> paramsList)
611-
throws ClientException {
612-
return requireAgreement(paramsList, OnlineQueryParamsComplete::getEnvironmentId, "environmentId");
621+
static @Nullable String resolveEnvironment(
622+
List<OnlineQueryParamsComplete> paramsList,
623+
@Nullable String clientEnvironmentId
624+
) throws ClientException {
625+
return requireAgreement(
626+
paramsList,
627+
params -> isBlank(params.getEnvironmentId())
628+
? clientEnvironmentId
629+
: params.getEnvironmentId(),
630+
"environmentId"
631+
);
613632
}
614633

615634
static @Nullable String resolveBranch(
@@ -623,11 +642,15 @@ static Optional<Duration> resolveDeadline(List<OnlineQueryParamsComplete> params
623642
);
624643
}
625644

645+
/*
646+
* Only the name has to agree. Envoy route-matches named queries on the
647+
* "x-chalk-query-name" header, of which the request has one. There is no
648+
* corresponding header for the version -- each sub-request carries its own in its
649+
* OnlineQueryContext -- so queries that differ only by version are satisfiable.
650+
*/
626651
static @Nullable String resolveQueryName(List<OnlineQueryParamsComplete> paramsList)
627652
throws ClientException {
628-
String queryName = requireAgreement(paramsList, OnlineQueryParamsComplete::getQueryName, "queryName");
629-
requireAgreement(paramsList, OnlineQueryParamsComplete::getQueryNameVersion, "queryNameVersion");
630-
return queryName;
653+
return requireAgreement(paramsList, OnlineQueryParamsComplete::getQueryName, "queryName");
631654
}
632655

633656
/*

‎src/main/java/ai/chalk/internal/arrow/Unmarshaller.java‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,17 @@ public class Unmarshaller {
3030
public static List<String> prefixToSkip = List.of(Constants.chalkDunderPrefix);
3131

3232
public static <T extends FeaturesClass> T[] unmarshalOnlineQueryResult(OnlineQueryResult result, Class<T> target) throws ClientException {
33+
if (result.getScalarsTable() == null) {
34+
// No scalar data came back. For a multi-query this is what a sub-query that failed
35+
// to execute leaves behind: an empty result whose reason sits in the multi-result's
36+
// global errors rather than on the result itself.
37+
throw new ClientException(
38+
"Cannot unmarshal an online query result that has no scalar data. Check "
39+
+ "getErrors() on this result, and, if it came from onlineQueryMulti, "
40+
+ "getGlobalErrors() on the OnlineQueryMultiResult -- a query that "
41+
+ "failed to execute is reported there."
42+
);
43+
}
3344
try {
3445
var rootFeatureClasses = unmarshalTable(result.getScalarsTable(), target);
3546
// This only unmarshals if has-many features gets returned in `result.groupsTables`

‎src/test/java/ai/chalk/client/TestGrpcMultiQuery.java‎

Lines changed: 103 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
import java.util.HashMap;
2525
import java.util.List;
2626
import java.util.Map;
27+
import java.util.Optional;
2728

2829
import static org.junit.jupiter.api.Assertions.assertEquals;
2930
import static org.junit.jupiter.api.Assertions.assertNotSame;
@@ -209,24 +210,57 @@ public void testBuildMultiRequestPreservesOrderAndUsesBulk() throws Exception {
209210

210211
// ------------------------------------------------------- settings resolvers
211212

212-
@Test
213-
public void testResolveDeadlineTakesTheLongest() {
214-
assertTrue(GRPCClient.resolveDeadline(List.of(simpleParams(), simpleParams())).isEmpty());
215-
216-
OnlineQueryParamsComplete shortTimeout = OnlineQueryParams.builder()
213+
private static OnlineQueryParamsComplete paramsWithTimeout(Duration timeout) {
214+
return OnlineQueryParams.builder()
217215
.withInput("user.id", List.of(1))
218216
.withOutputs("user.socure_score")
219-
.withTimeout(Duration.ofSeconds(1))
220-
.build();
221-
OnlineQueryParamsComplete longTimeout = OnlineQueryParams.builder()
222-
.withInput("user.id", List.of(1))
223-
.withOutputs("user.socure_score")
224-
.withTimeout(Duration.ofSeconds(30))
217+
.withTimeout(timeout)
225218
.build();
219+
}
220+
221+
@Test
222+
public void testResolveDeadlineTakesTheLongest() {
223+
var noClientTimeout = Optional.<Duration>empty();
224+
225+
// nothing set anywhere -> no deadline
226+
assertTrue(GRPCClient.resolveDeadline(
227+
List.of(simpleParams(), simpleParams()), noClientTimeout).isEmpty());
228+
229+
var shortTimeout = paramsWithTimeout(Duration.ofSeconds(1));
230+
var longTimeout = paramsWithTimeout(Duration.ofSeconds(30));
231+
232+
assertEquals(
233+
Duration.ofSeconds(30),
234+
GRPCClient.resolveDeadline(List.of(shortTimeout, longTimeout), noClientTimeout)
235+
.orElseThrow());
236+
}
237+
238+
/**
239+
* A query that sets no timeout is bounded by the client-level timeout, so that is
240+
* its effective timeout. Ignoring it would let one query's short timeout cut short
241+
* a query that was relying on a longer client default.
242+
*/
243+
@Test
244+
public void testResolveDeadlineHonoursClientLevelTimeout() {
245+
var fast = paramsWithTimeout(Duration.ofSeconds(1));
246+
var clientTimeout = Optional.of(Duration.ofMinutes(5));
247+
248+
// `fast` wants 1s, the unset one is entitled to the client's 5m -> 5m wins
249+
assertEquals(
250+
Duration.ofMinutes(5),
251+
GRPCClient.resolveDeadline(List.of(fast, simpleParams()), clientTimeout).orElseThrow());
226252

253+
// every query explicit -> the client default is irrelevant
227254
assertEquals(
228255
Duration.ofSeconds(30),
229-
GRPCClient.resolveDeadline(List.of(shortTimeout, longTimeout, simpleParams())).orElseThrow());
256+
GRPCClient.resolveDeadline(
257+
List.of(fast, paramsWithTimeout(Duration.ofSeconds(30))), clientTimeout)
258+
.orElseThrow());
259+
260+
// an unset query with no client timeout is unbounded, so the request cannot
261+
// carry a deadline at all
262+
assertTrue(GRPCClient.resolveDeadline(
263+
List.of(fast, simpleParams()), Optional.empty()).isEmpty());
230264
}
231265

232266
@Test
@@ -246,51 +280,58 @@ public void testResolveQueryNameRequiresAgreement() throws Exception {
246280
() -> GRPCClient.resolveQueryName(List.of(params("a"), simpleParams())));
247281
}
248282

249-
/**
250-
* When both the name and the version disagree, the name is the more useful
251-
* thing to report: it is what the caller actually set.
252-
*/
253283
@Test
254-
public void testResolveQueryNameReportsNameBeforeVersion() {
255-
OnlineQueryParamsComplete a = OnlineQueryParams.builder()
284+
public void testResolveEnvironmentRequiresAgreement() throws Exception {
285+
OnlineQueryParamsComplete envA = OnlineQueryParams.builder()
256286
.withInput("user.id", List.of(1))
257287
.withOutputs("user.socure_score")
258-
.withQueryName("a")
259-
.withQueryNameVersion("1")
288+
.withEnvironmentId("env-a")
260289
.build();
261-
OnlineQueryParamsComplete b = OnlineQueryParams.builder()
290+
OnlineQueryParamsComplete envB = OnlineQueryParams.builder()
262291
.withInput("user.id", List.of(1))
263292
.withOutputs("user.socure_score")
264-
.withQueryName("b")
265-
.withQueryNameVersion("2")
293+
.withEnvironmentId("env-b")
266294
.build();
267295

296+
assertEquals("env-a", GRPCClient.resolveEnvironment(List.of(envA, envA), null));
297+
268298
var conflict = assertThrows(
269-
ClientException.class, () -> GRPCClient.resolveQueryName(List.of(a, b)));
270-
assertTrue(conflict.getMessage().contains("queryName,"));
271-
assertTrue(!conflict.getMessage().contains("queryNameVersion"));
299+
ClientException.class,
300+
() -> GRPCClient.resolveEnvironment(List.of(envA, envB), null));
301+
assertTrue(conflict.getMessage().contains("environmentId"));
302+
assertTrue(conflict.getMessage().contains("'env-a'"));
303+
assertTrue(conflict.getMessage().contains("'env-b'"));
304+
305+
// An unset environmentId means the client's environment, so naming that same
306+
// environment explicitly on one query is not a conflict.
307+
assertEquals(
308+
"env-a",
309+
GRPCClient.resolveEnvironment(List.of(envA, simpleParams()), "env-a"));
310+
assertThrows(
311+
ClientException.class,
312+
() -> GRPCClient.resolveEnvironment(List.of(envA, simpleParams()), "env-other"));
272313
}
273314

315+
/**
316+
* The version rides along in each sub-request's own context and has no header of
317+
* its own, so queries that differ only by version are satisfiable.
318+
*/
274319
@Test
275-
public void testResolveEnvironmentRequiresAgreement() throws Exception {
276-
OnlineQueryParamsComplete envA = OnlineQueryParams.builder()
320+
public void testResolveQueryNameAllowsDifferingVersions() throws Exception {
321+
OnlineQueryParamsComplete v1 = OnlineQueryParams.builder()
277322
.withInput("user.id", List.of(1))
278323
.withOutputs("user.socure_score")
279-
.withEnvironmentId("env-a")
324+
.withQueryName("same")
325+
.withQueryNameVersion("1")
280326
.build();
281-
OnlineQueryParamsComplete envB = OnlineQueryParams.builder()
327+
OnlineQueryParamsComplete v2 = OnlineQueryParams.builder()
282328
.withInput("user.id", List.of(1))
283329
.withOutputs("user.socure_score")
284-
.withEnvironmentId("env-b")
330+
.withQueryName("same")
331+
.withQueryNameVersion("2")
285332
.build();
286333

287-
assertEquals("env-a", GRPCClient.resolveEnvironment(List.of(envA, envA)));
288-
289-
var conflict = assertThrows(
290-
ClientException.class, () -> GRPCClient.resolveEnvironment(List.of(envA, envB)));
291-
assertTrue(conflict.getMessage().contains("environmentId"));
292-
assertTrue(conflict.getMessage().contains("'env-a'"));
293-
assertTrue(conflict.getMessage().contains("'env-b'"));
334+
assertEquals("same", GRPCClient.resolveQueryName(List.of(v1, v2)));
294335
}
295336

296337
@Test
@@ -441,6 +482,30 @@ public void testDecodeMultiResponseExecutionFailureIsGlobalWithEmptySlot() throw
441482
assertEquals(0, this.allocator.getAllocatedMemory());
442483
}
443484

485+
/**
486+
* Unmarshalling the empty slot left by a query that failed to execute must explain
487+
* itself rather than surfacing an opaque NPE from deep inside the Arrow reader.
488+
*/
489+
@Test
490+
public void testUnmarshalEmptySlotExplainsItself() throws Exception {
491+
var response = OnlineQueryMultiResponse.newBuilder()
492+
.addResponses(wrap(OnlineQueryBulkResponse.newBuilder().build()))
493+
.addErrors(error("Encountered an error occurred while executing query 1/1: rate limited"))
494+
.build();
495+
496+
try (OnlineQueryMultiResult multi =
497+
GRPCClient.decodeMultiResponse(response, 1, TRACE_ID, this.allocator)) {
498+
OnlineQueryResult empty = multi.getResults().get(0);
499+
assertNull(empty.getScalarsTable());
500+
501+
var thrown = assertThrows(
502+
ClientException.class, () -> empty.unmarshal(TestGrpcMultiQueryUser.class));
503+
assertTrue(thrown.getMessage().contains("no scalar data"));
504+
assertTrue(thrown.getMessage().contains("getGlobalErrors()"));
505+
}
506+
assertEquals(0, this.allocator.getAllocatedMemory());
507+
}
508+
444509
@Test
445510
public void testDecodeMultiResponseCountMismatch() throws Exception {
446511
var tooFew = OnlineQueryMultiResponse.newBuilder()
Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
package ai.chalk.client;
2+
3+
import ai.chalk.features.Feature;
4+
import ai.chalk.features.FeaturesClass;
5+
6+
/** Minimal feature class used only to exercise the unmarshal guard. */
7+
public class TestGrpcMultiQueryUser extends FeaturesClass {
8+
public Feature<String> id;
9+
public Feature<Double> socure_score;
10+
}

0 commit comments

Comments
 (0)