From 7c6f859dc521cc0dbeac6ba7624d67cb82d4b444 Mon Sep 17 00:00:00 2001 From: t Date: Mon, 31 Aug 2026 08:07:31 +0300 Subject: [PATCH 1/6] [OPIK-8175] [BE] perf: skip the dataset item-count scan when a version supplies items_total enrichDatasetWithAdditionalInformation ran an O(N) count(DISTINCT id) over dataset_items unconditionally, then discarded the result whenever dataset versioning supplied itemsTotal. Move the latest-version lookup ahead of the count and narrow the count query to the datasets that actually need the legacy fallback, skipping it entirely when none do. Co-Authored-By: Claude Opus 5 (1M context) --- .../com/comet/opik/domain/DatasetService.java | 50 +++- .../domain/DatasetServiceEnrichmentTest.java | 257 ++++++++++++++++++ 2 files changed, 292 insertions(+), 15 deletions(-) create mode 100644 apps/opik-backend/src/test/java/com/comet/opik/domain/DatasetServiceEnrichmentTest.java diff --git a/apps/opik-backend/src/main/java/com/comet/opik/domain/DatasetService.java b/apps/opik-backend/src/main/java/com/comet/opik/domain/DatasetService.java index 53b6ea358b7..9e09e22602f 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/domain/DatasetService.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/domain/DatasetService.java @@ -713,11 +713,6 @@ private List enrichDatasetWithAdditionalInformation(List datas .toStream() .collect(toMap(ExperimentSummary::datasetId, Function.identity())); - Map datasetItemSummaryMap = datasetItemDAO.findDatasetItemSummaryByDatasetIds(ids) - .contextWrite(ctx -> AsyncUtils.setRequestContext(ctx, requestContext)) - .toStream() - .collect(toMap(DatasetItemSummary::datasetId, Function.identity())); - Map optimizationSummaryMap = optimizationDAO .findOptimizationSummaryByDatasetIds(ids) .contextWrite(ctx -> AsyncUtils.setRequestContext(ctx, requestContext)) @@ -726,23 +721,29 @@ private List enrichDatasetWithAdditionalInformation(List datas Map latestVersionsByDatasetId = fetchLatestVersionsByDatasetIds(ids); + // The dataset_items count is O(N) in the dataset's item count, so only pay it for the datasets whose + // count cannot come from their latest version. + Set idsNeedingItemCount = datasets.stream() + .map(Dataset::id) + .filter(id -> versionItemsTotal(latestVersionsByDatasetId.get(id)) == null) + .collect(toSet()); + + Map datasetItemSummaryMap = fetchDatasetItemSummaries(idsNeedingItemCount); + return datasets.stream() .map(dataset -> { var resume = experimentSummary.computeIfAbsent(dataset.id(), ExperimentSummary::empty); - var datasetItemSummary = datasetItemSummaryMap.computeIfAbsent(dataset.id(), - DatasetItemSummary::empty); var optimizationSummary = optimizationSummaryMap.computeIfAbsent(dataset.id(), OptimizationDAO.OptimizationSummary::empty); var latestVersion = latestVersionsByDatasetId.get(dataset.id()); - // When versioning is enabled and a latest version exists, use itemsTotal from the version - // Otherwise, fall back to the legacy dataset_items count - Long itemsCount; - if (featureFlags.isDatasetVersioningEnabled() && latestVersion != null - && latestVersion.itemsTotal() != null) { - itemsCount = latestVersion.itemsTotal().longValue(); - } else { - itemsCount = datasetItemSummary.datasetItemsCount(); + // When versioning is enabled and a latest version supplies itemsTotal, use it. + // Otherwise, fall back to the legacy dataset_items count. + Long itemsCount = versionItemsTotal(latestVersion); + if (itemsCount == null) { + itemsCount = datasetItemSummaryMap + .getOrDefault(dataset.id(), DatasetItemSummary.empty(dataset.id())) + .datasetItemsCount(); } return dataset.toBuilder() @@ -757,6 +758,25 @@ private List enrichDatasetWithAdditionalInformation(List datas .toList(); } + private Long versionItemsTotal(DatasetVersion latestVersion) { + if (!featureFlags.isDatasetVersioningEnabled() || latestVersion == null + || latestVersion.itemsTotal() == null) { + return null; + } + return latestVersion.itemsTotal().longValue(); + } + + private Map fetchDatasetItemSummaries(Set datasetIds) { + if (datasetIds.isEmpty()) { + return Map.of(); + } + + return datasetItemDAO.findDatasetItemSummaryByDatasetIds(datasetIds) + .contextWrite(ctx -> AsyncUtils.setRequestContext(ctx, requestContext)) + .toStream() + .collect(toMap(DatasetItemSummary::datasetId, Function.identity())); + } + private Map fetchLatestVersionsByDatasetIds(Set datasetIds) { if (datasetIds.isEmpty()) { return Map.of(); diff --git a/apps/opik-backend/src/test/java/com/comet/opik/domain/DatasetServiceEnrichmentTest.java b/apps/opik-backend/src/test/java/com/comet/opik/domain/DatasetServiceEnrichmentTest.java new file mode 100644 index 00000000000..15e419987a5 --- /dev/null +++ b/apps/opik-backend/src/test/java/com/comet/opik/domain/DatasetServiceEnrichmentTest.java @@ -0,0 +1,257 @@ +package com.comet.opik.domain; + +import com.comet.opik.api.Dataset; +import com.comet.opik.api.DatasetVersion; +import com.comet.opik.api.sorting.SortingFactoryDatasets; +import com.comet.opik.domain.filter.FilterQueryBuilder; +import com.comet.opik.domain.sorting.SortingQueryBuilder; +import com.comet.opik.infrastructure.BatchOperationsConfig; +import com.comet.opik.infrastructure.FeatureFlags; +import com.comet.opik.infrastructure.auth.RequestContext; +import com.google.common.eventbus.EventBus; +import org.jdbi.v3.core.Handle; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; +import org.mockito.junit.jupiter.MockitoSettings; +import org.mockito.quality.Strictness; +import reactor.core.publisher.Flux; +import ru.vyarus.guicey.jdbi3.tx.TransactionTemplate; +import ru.vyarus.guicey.jdbi3.tx.TxAction; + +import java.util.Arrays; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.UUID; +import java.util.stream.Collectors; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +/** + * Covers the item-count source selection in {@code enrichDatasetWithAdditionalInformation}: the legacy + * {@code dataset_items} count is an O(N) scan, so it must only be issued for datasets that cannot take their + * count from a dataset version. + */ +@ExtendWith(MockitoExtension.class) +@MockitoSettings(strictness = Strictness.LENIENT) +class DatasetServiceEnrichmentTest { + + private static final String WORKSPACE_ID = "test-workspace"; + private static final String WORKSPACE_NAME = "test-workspace-name"; + private static final String USER_NAME = "test-user"; + + @Mock + private IdGenerator idGenerator; + @Mock + private TransactionTemplate template; + @Mock + private RequestContext requestContext; + @Mock + private ExperimentItemDAO experimentItemDAO; + @Mock + private DatasetItemDAO datasetItemDAO; + @Mock + private ExperimentDAO experimentDAO; + @Mock + private SortingQueryBuilder sortingQueryBuilder; + @Mock + private FilterQueryBuilder filterQueryBuilder; + @Mock + private SortingFactoryDatasets sortingFactory; + @Mock + private BatchOperationsConfig batchOperationsConfig; + @Mock + private OptimizationDAO optimizationDAO; + @Mock + private EventBus eventBus; + @Mock + private FeatureFlags featureFlags; + @Mock + private ProjectService projectService; + @Mock + private Handle handle; + @Mock + private DatasetDAO datasetDAO; + @Mock + private DatasetVersionDAO datasetVersionDAO; + + private DatasetServiceImpl service; + + @BeforeEach + void setUp() { + service = new DatasetServiceImpl(idGenerator, template, () -> requestContext, experimentItemDAO, + datasetItemDAO, experimentDAO, sortingQueryBuilder, filterQueryBuilder, sortingFactory, + batchOperationsConfig, optimizationDAO, eventBus, featureFlags, projectService); + + when(requestContext.getWorkspaceId()).thenReturn(WORKSPACE_ID); + when(requestContext.getWorkspaceName()).thenReturn(WORKSPACE_NAME); + when(requestContext.getUserName()).thenReturn(USER_NAME); + + when(template.inTransaction(any(), any())).thenAnswer(invocation -> { + TxAction callback = invocation.getArgument(1); + return callback.execute(handle); + }); + when(handle.attach(DatasetDAO.class)).thenReturn(datasetDAO); + when(handle.attach(DatasetVersionDAO.class)).thenReturn(datasetVersionDAO); + + when(experimentItemDAO.findExperimentSummaryByDatasetIds(any())).thenReturn(Flux.empty()); + when(optimizationDAO.findOptimizationSummaryByDatasetIds(any())).thenReturn(Flux.empty()); + when(datasetItemDAO.findDatasetItemSummaryByDatasetIds(any())).thenReturn(Flux.empty()); + + when(featureFlags.isDatasetVersioningEnabled()).thenReturn(true); + when(sortingFactory.getSortableFields()).thenReturn(List.of()); + } + + @Test + @DisplayName("no dataset_items count is issued when every dataset resolves its count from a version") + void fullyVersionedBatchSkipsItemCountQuery() { + var first = UUID.randomUUID(); + var second = UUID.randomUUID(); + givenDatasets(first, second); + givenLatestVersions(version(first, 100), version(second, 250)); + + var content = findPage().content(); + + verify(datasetItemDAO, never()).findDatasetItemSummaryByDatasetIds(any()); + assertThat(itemsCountById(content)).containsExactlyInAnyOrderEntriesOf(Map.of(first, 100L, second, 250L)); + } + + @Test + @DisplayName("dataset_items count is issued and used when versioning is disabled") + void versioningDisabledFallsBackToItemCount() { + when(featureFlags.isDatasetVersioningEnabled()).thenReturn(false); + + var id = UUID.randomUUID(); + givenDatasets(id); + givenLatestVersions(version(id, 100)); + givenItemCounts(Map.of(id, 7L)); + + var content = findPage().content(); + + verify(datasetItemDAO).findDatasetItemSummaryByDatasetIds(Set.of(id)); + assertThat(itemsCountById(content)).containsExactlyInAnyOrderEntriesOf(Map.of(id, 7L)); + } + + @Test + @DisplayName("dataset_items count is issued and used when no latest version exists") + void missingVersionFallsBackToItemCount() { + var id = UUID.randomUUID(); + givenDatasets(id); + givenLatestVersions(); + givenItemCounts(Map.of(id, 42L)); + + var content = findPage().content(); + + verify(datasetItemDAO).findDatasetItemSummaryByDatasetIds(Set.of(id)); + assertThat(itemsCountById(content)).containsExactlyInAnyOrderEntriesOf(Map.of(id, 42L)); + } + + @Test + @DisplayName("dataset_items count is issued and used when the latest version has a null itemsTotal") + void nullItemsTotalFallsBackToItemCount() { + var id = UUID.randomUUID(); + givenDatasets(id); + givenLatestVersions(version(id, null)); + givenItemCounts(Map.of(id, 13L)); + + var content = findPage().content(); + + verify(datasetItemDAO).findDatasetItemSummaryByDatasetIds(Set.of(id)); + assertThat(itemsCountById(content)).containsExactlyInAnyOrderEntriesOf(Map.of(id, 13L)); + } + + @Test + @DisplayName("a dataset with no version and no items reports a zero count") + void missingVersionAndMissingItemCountYieldsZero() { + var id = UUID.randomUUID(); + givenDatasets(id); + givenLatestVersions(); + givenItemCounts(Map.of()); + + var content = findPage().content(); + + assertThat(itemsCountById(content)).containsExactlyInAnyOrderEntriesOf(Map.of(id, 0L)); + } + + @Test + @DisplayName("a mixed batch resolves every count correctly in a single narrowed item-count query") + void mixedBatchNarrowsItemCountQueryToTheFallbackSubset() { + var versioned = UUID.randomUUID(); + var noVersion = UUID.randomUUID(); + var nullTotal = UUID.randomUUID(); + givenDatasets(versioned, noVersion, nullTotal); + givenLatestVersions(version(versioned, 500), version(nullTotal, null)); + givenItemCounts(Map.of(noVersion, 3L, nullTotal, 9L)); + + var content = findPage().content(); + + verify(datasetItemDAO).findDatasetItemSummaryByDatasetIds(Set.of(noVersion, nullTotal)); + assertThat(itemsCountById(content)).containsExactlyInAnyOrderEntriesOf( + Map.of(versioned, 500L, noVersion, 3L, nullTotal, 9L)); + } + + @Test + @DisplayName("the single-dataset retrieve path also skips the item-count query for a versioned dataset") + void findByIdSkipsItemCountForVersionedDataset() { + var id = UUID.randomUUID(); + var dataset = dataset(id); + when(datasetDAO.findById(id, WORKSPACE_ID)).thenReturn(Optional.of(dataset)); + givenLatestVersions(version(id, 77)); + + var result = service.findById(id); + + verify(datasetItemDAO, never()).findDatasetItemSummaryByDatasetIds(any()); + assertThat(result.datasetItemsCount()).isEqualTo(77L); + } + + private Dataset.DatasetPage findPage() { + return service.find(1, 10, DatasetCriteria.builder().build(), List.of()); + } + + private void givenDatasets(UUID... ids) { + List datasets = Arrays.stream(ids).map(this::dataset).toList(); + when(datasetDAO.find(eq(10), eq(0), eq(WORKSPACE_ID), any(), any(), eq(false), eq(false), any(), any(), any(), + any())).thenReturn(datasets); + when(datasetDAO.findCount(eq(WORKSPACE_ID), any(), any(), eq(false), eq(false), any(), any(), any())) + .thenReturn((long) ids.length); + } + + private void givenLatestVersions(DatasetVersion... versions) { + when(datasetVersionDAO.findLatestVersionsByDatasetIds(any(), anyString())).thenReturn(List.of(versions)); + } + + private void givenItemCounts(Map countsByDatasetId) { + var summaries = countsByDatasetId.entrySet().stream() + .map(entry -> new DatasetItemSummary(entry.getKey(), entry.getValue())) + .toList(); + when(datasetItemDAO.findDatasetItemSummaryByDatasetIds(any())).thenReturn(Flux.fromIterable(summaries)); + } + + private Dataset dataset(UUID id) { + return Dataset.builder().id(id).name("dataset-" + id).build(); + } + + private DatasetVersion version(UUID datasetId, Integer itemsTotal) { + return DatasetVersion.builder() + .id(UUID.randomUUID()) + .datasetId(datasetId) + .itemsTotal(itemsTotal) + .isLatest(true) + .build(); + } + + private Map itemsCountById(List datasets) { + return datasets.stream().collect(Collectors.toMap(Dataset::id, Dataset::datasetItemsCount)); + } +} From 1eabc6581e9c8586ba485b60cc4403dcedcfcb5f Mon Sep 17 00:00:00 2001 From: t Date: Mon, 31 Aug 2026 08:33:59 +0300 Subject: [PATCH 2/6] fix(datasets): treat the not-migrated sentinel as an unavailable item count versionItemsTotal accepted any non-null itemsTotal as authoritative, including the ITEMS_TOTAL_NOT_MIGRATED (-1) sentinel written by migration 000046 for versions awaiting backfill. main had the same hole, but narrowing the fallback made it reachable: the dataset was excluded from the count query entirely, so -1 became the only possible answer instead of merely a wasted query. Also drop the eager DatasetItemSummary.empty allocation that getOrDefault evaluated on every dataset regardless of whether the map had a match. Co-Authored-By: Claude Opus 5 (1M context) --- .../java/com/comet/opik/domain/DatasetService.java | 10 ++++++---- .../opik/domain/DatasetServiceEnrichmentTest.java | 14 ++++++++++++++ 2 files changed, 20 insertions(+), 4 deletions(-) diff --git a/apps/opik-backend/src/main/java/com/comet/opik/domain/DatasetService.java b/apps/opik-backend/src/main/java/com/comet/opik/domain/DatasetService.java index 9e09e22602f..c3d32cd963d 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/domain/DatasetService.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/domain/DatasetService.java @@ -741,9 +741,10 @@ private List enrichDatasetWithAdditionalInformation(List datas // Otherwise, fall back to the legacy dataset_items count. Long itemsCount = versionItemsTotal(latestVersion); if (itemsCount == null) { - itemsCount = datasetItemSummaryMap - .getOrDefault(dataset.id(), DatasetItemSummary.empty(dataset.id())) - .datasetItemsCount(); + var datasetItemSummary = datasetItemSummaryMap.get(dataset.id()); + itemsCount = datasetItemSummary != null + ? datasetItemSummary.datasetItemsCount() + : DatasetItemSummary.empty(dataset.id()).datasetItemsCount(); } return dataset.toBuilder() @@ -760,7 +761,8 @@ private List enrichDatasetWithAdditionalInformation(List datas private Long versionItemsTotal(DatasetVersion latestVersion) { if (!featureFlags.isDatasetVersioningEnabled() || latestVersion == null - || latestVersion.itemsTotal() == null) { + || latestVersion.itemsTotal() == null + || latestVersion.itemsTotal() == DatasetVersionDAO.ITEMS_TOTAL_NOT_MIGRATED) { return null; } return latestVersion.itemsTotal().longValue(); diff --git a/apps/opik-backend/src/test/java/com/comet/opik/domain/DatasetServiceEnrichmentTest.java b/apps/opik-backend/src/test/java/com/comet/opik/domain/DatasetServiceEnrichmentTest.java index 15e419987a5..060d82a56af 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/domain/DatasetServiceEnrichmentTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/domain/DatasetServiceEnrichmentTest.java @@ -171,6 +171,20 @@ void nullItemsTotalFallsBackToItemCount() { assertThat(itemsCountById(content)).containsExactlyInAnyOrderEntriesOf(Map.of(id, 13L)); } + @Test + @DisplayName("the not-migrated sentinel is not an authoritative count and falls back to dataset_items") + void notMigratedSentinelFallsBackToItemCount() { + var id = UUID.randomUUID(); + givenDatasets(id); + givenLatestVersions(version(id, DatasetVersionDAO.ITEMS_TOTAL_NOT_MIGRATED)); + givenItemCounts(Map.of(id, 21L)); + + var content = findPage().content(); + + verify(datasetItemDAO).findDatasetItemSummaryByDatasetIds(Set.of(id)); + assertThat(itemsCountById(content)).containsExactlyInAnyOrderEntriesOf(Map.of(id, 21L)); + } + @Test @DisplayName("a dataset with no version and no items reports a zero count") void missingVersionAndMissingItemCountYieldsZero() { From ff1fff0ba487f27919fd2d8c15f99fba398a05db Mon Sep 17 00:00:00 2001 From: Daniel Dimenshtein Date: Thu, 3 Sep 2026 09:14:42 +0300 Subject: [PATCH 3/6] docs(datasets): correct the enrichment zip comment after the merge resolution Co-Authored-By: Claude Opus 5 (1M context) --- .../src/main/java/com/comet/opik/domain/DatasetService.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/apps/opik-backend/src/main/java/com/comet/opik/domain/DatasetService.java b/apps/opik-backend/src/main/java/com/comet/opik/domain/DatasetService.java index c8c16f9b394..9b014d8b1d6 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/domain/DatasetService.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/domain/DatasetService.java @@ -728,7 +728,7 @@ private List enrichDatasetWithAdditionalInformation(List datas // The dataset_items count is deliberately NOT part of this zip: it is O(N) in each dataset's item count, // and which datasets need it is only known once the latest versions are in hand. It is chained below - // instead, over the narrowed id set. The other three stay concurrent. + // instead, over the narrowed id set. The remaining lookups stay concurrent. // // collect(...) with Collectors.toMap rather than Flux.collectMap: collectMap is last-wins, whereas the // serial code this replaces threw on a duplicate dataset_id. All queries GROUP BY dataset_id so From 1e4b884f5308a2fc61e38b57cad4515082bc0a29 Mon Sep 17 00:00:00 2001 From: Daniel Dimenshtein Date: Thu, 3 Sep 2026 09:37:47 +0300 Subject: [PATCH 4/6] perf(datasets): start the fallback item count from the version lookup alone Chaining fetchDatasetItemSummaries off the whole Mono.zip made it wait on the experiment and optimization summaries, which it has no dependency on: the zip only emits once all three sources complete, so a fallback-heavy page serialized a ClickHouse round trip that could have overlapped. cache() the version lookup so one execution feeds both the zip and the count chain, then zip the chain alongside the other two lookups. The count is issued as soon as the versions resolve, and the narrowed id set is unchanged. Co-Authored-By: Claude Opus 5 (1M context) --- .../com/comet/opik/domain/DatasetService.java | 28 +++++---- .../domain/DatasetServiceEnrichmentTest.java | 57 +++++++++++++++++++ 2 files changed, 74 insertions(+), 11 deletions(-) diff --git a/apps/opik-backend/src/main/java/com/comet/opik/domain/DatasetService.java b/apps/opik-backend/src/main/java/com/comet/opik/domain/DatasetService.java index 9b014d8b1d6..e93369aa1bf 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/domain/DatasetService.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/domain/DatasetService.java @@ -41,7 +41,6 @@ import org.jdbi.v3.core.statement.UnableToExecuteStatementException; import reactor.core.publisher.Mono; import reactor.core.scheduler.Schedulers; -import reactor.util.function.Tuples; import ru.vyarus.dropwizard.guice.module.yaml.bind.Config; import ru.vyarus.guicey.jdbi3.tx.TransactionTemplate; @@ -726,26 +725,33 @@ private List enrichDatasetWithAdditionalInformation(List datas String workspaceId = requestContext.get().getWorkspaceId(); String userName = requestContext.get().getUserName(); - // The dataset_items count is deliberately NOT part of this zip: it is O(N) in each dataset's item count, - // and which datasets need it is only known once the latest versions are in hand. It is chained below - // instead, over the narrowed id set. The remaining lookups stay concurrent. + // The dataset_items count cannot join the zip directly: it is O(N) in each dataset's item count, and + // which datasets need it is only known once the latest versions are in hand. So it is chained off the + // version lookup alone -- not off the whole zip, which would make it wait on the experiment and + // optimization summaries it has no dependency on and serialize a round trip on fallback-heavy pages. + // cache() lets the version result feed both the chain and the zip from one execution. // // collect(...) with Collectors.toMap rather than Flux.collectMap: collectMap is last-wins, whereas the // serial code this replaces threw on a duplicate dataset_id. All queries GROUP BY dataset_id so // duplicates should not occur; keeping the loud form means a query change that broke that assumption // fails instead of silently dropping one row's summary. + // + // defaultIfEmpty guards the zip: a Mono that completes empty makes zip emit nothing at all, which + // would turn an absent-versions result into a null and NPE below. + Mono> latestVersions = Mono + .fromCallable(() -> fetchLatestVersionsByDatasetIds(ids, workspaceId)) + .subscribeOn(Schedulers.boundedElastic()) + .defaultIfEmpty(Map.of()) + .cache(); + var enrichmentData = Mono.zip( experimentItemDAO.findExperimentSummaryByDatasetIds(ids) .collect(toMap(ExperimentSummary::datasetId, Function.identity())), optimizationDAO.findOptimizationSummaryByDatasetIds(ids) .collect(toMap(OptimizationDAO.OptimizationSummary::datasetId, Function.identity())), - // defaultIfEmpty guards the zip: a Mono that completes empty makes zip emit nothing at all, - // which would turn an absent-versions result into a null and NPE below. - Mono.fromCallable(() -> fetchLatestVersionsByDatasetIds(ids, workspaceId)) - .subscribeOn(Schedulers.boundedElastic()) - .defaultIfEmpty(Map.of())) - .flatMap(data -> fetchDatasetItemSummaries(idsNeedingItemCount(datasets, data.getT3())) - .map(itemSummaries -> Tuples.of(data.getT1(), data.getT2(), data.getT3(), itemSummaries))) + latestVersions, + latestVersions.flatMap( + versions -> fetchDatasetItemSummaries(idsNeedingItemCount(datasets, versions)))) .contextWrite(ctx -> AsyncUtils.setRequestContext(ctx, userName, workspaceId)) .block(); diff --git a/apps/opik-backend/src/test/java/com/comet/opik/domain/DatasetServiceEnrichmentTest.java b/apps/opik-backend/src/test/java/com/comet/opik/domain/DatasetServiceEnrichmentTest.java index 2e457a71b0d..393810ff951 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/domain/DatasetServiceEnrichmentTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/domain/DatasetServiceEnrichmentTest.java @@ -625,6 +625,63 @@ void enrichmentDefersItemCountUntilVersionsAreKnown() { assertThat(itemsCountById(actual.content())).containsExactlyInAnyOrderEntriesOf(Map.of(datasetId, 4L)); } + @Test + @DisplayName("The item-count query does not wait on the experiment and optimization summaries") + void enrichmentDoesNotBlockItemCountOnUnrelatedSummaries() throws Exception { + var datasetId = UUID.randomUUID(); + stubPage(datasetId); + stubVersioningEnabled(); + stubLatestVersions(); + + // The item count depends only on the versions, so it must be subscribed while the two unrelated + // summaries are still in flight. Both are held open until the count has been issued; if the count + // were chained off the whole zip it would never be reached and the latch would time out. + var itemCountSubscribed = new CountDownLatch(1); + var releaseSummaries = new CountDownLatch(1); + + Flux heldExperiments = Flux.defer(() -> { + awaitQuietly(releaseSummaries); + return Flux.empty(); + }).subscribeOn(Schedulers.boundedElastic()); + + Flux heldOptimizations = Flux.defer(() -> { + awaitQuietly(releaseSummaries); + return Flux.empty(); + }).subscribeOn(Schedulers.boundedElastic()); + + when(experimentItemDAO.findExperimentSummaryByDatasetIds(anySet())).thenReturn(heldExperiments); + when(optimizationDAO.findOptimizationSummaryByDatasetIds(anySet())).thenReturn(heldOptimizations); + when(datasetItemDAO.findDatasetItemSummaryByDatasetIds(anySet())).thenAnswer(invocation -> { + itemCountSubscribed.countDown(); + return Flux.just(new DatasetItemSummary(datasetId, 6)); + }); + + var enrichment = CompletableFuture + .supplyAsync(() -> service.find(1, 10, DatasetCriteria.builder().build(), List.of())); + + try { + assertThat(itemCountSubscribed.await(SUBSCRIBE_TIMEOUT_SECONDS, TimeUnit.SECONDS)) + .as("the item-count query should be issued while the unrelated summaries are still pending") + .isTrue(); + } finally { + releaseSummaries.countDown(); + } + + var actual = enrichment.get(WORKER_RELEASE_TIMEOUT_SECONDS, TimeUnit.SECONDS); + assertThat(itemsCountById(actual.content())).containsExactlyInAnyOrderEntriesOf(Map.of(datasetId, 6L)); + } + + private static void awaitQuietly(CountDownLatch latch) { + try { + if (!latch.await(WORKER_RELEASE_TIMEOUT_SECONDS, TimeUnit.SECONDS)) { + throw new IllegalStateException("summaries were never released"); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IllegalStateException(e); + } + } + private void stubVersioningEnabled() { when(featureFlags.isDatasetVersioningEnabled()).thenReturn(true); } From e30a82b01eac48ed9dbea28ab803a57b2e3d3a40 Mon Sep 17 00:00:00 2001 From: Daniel Dimenshtein Date: Thu, 3 Sep 2026 10:33:26 +0300 Subject: [PATCH 5/6] test(datasets): name the item-count latch for what it observes The Mockito answer fires when findDatasetItemSummaryByDatasetIds is invoked, not when the returned Flux is subscribed, so itemCountSubscribed overstated the synchronization point. Rename to itemCountQueryIssued and spell out in the comment that invocation -- the point the narrowed id set is handed over -- is the property under test. Co-Authored-By: Claude Opus 5 (1M context) --- .../opik/domain/DatasetServiceEnrichmentTest.java | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/apps/opik-backend/src/test/java/com/comet/opik/domain/DatasetServiceEnrichmentTest.java b/apps/opik-backend/src/test/java/com/comet/opik/domain/DatasetServiceEnrichmentTest.java index 393810ff951..60d15bf98e5 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/domain/DatasetServiceEnrichmentTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/domain/DatasetServiceEnrichmentTest.java @@ -633,10 +633,12 @@ void enrichmentDoesNotBlockItemCountOnUnrelatedSummaries() throws Exception { stubVersioningEnabled(); stubLatestVersions(); - // The item count depends only on the versions, so it must be subscribed while the two unrelated - // summaries are still in flight. Both are held open until the count has been issued; if the count - // were chained off the whole zip it would never be reached and the latch would time out. - var itemCountSubscribed = new CountDownLatch(1); + // The item count depends only on the versions, so the query must be issued while the two unrelated + // summaries are still in flight. The latch counts down in the Mockito answer -- i.e. when the DAO + // method is invoked, which is the point the narrowed id set is computed and handed over, not the + // later subscription to the returned Flux. That is the property under test: if the count were + // chained off the whole zip it would never be requested and the latch would time out. + var itemCountQueryIssued = new CountDownLatch(1); var releaseSummaries = new CountDownLatch(1); Flux heldExperiments = Flux.defer(() -> { @@ -652,7 +654,7 @@ void enrichmentDoesNotBlockItemCountOnUnrelatedSummaries() throws Exception { when(experimentItemDAO.findExperimentSummaryByDatasetIds(anySet())).thenReturn(heldExperiments); when(optimizationDAO.findOptimizationSummaryByDatasetIds(anySet())).thenReturn(heldOptimizations); when(datasetItemDAO.findDatasetItemSummaryByDatasetIds(anySet())).thenAnswer(invocation -> { - itemCountSubscribed.countDown(); + itemCountQueryIssued.countDown(); return Flux.just(new DatasetItemSummary(datasetId, 6)); }); @@ -660,7 +662,7 @@ void enrichmentDoesNotBlockItemCountOnUnrelatedSummaries() throws Exception { .supplyAsync(() -> service.find(1, 10, DatasetCriteria.builder().build(), List.of())); try { - assertThat(itemCountSubscribed.await(SUBSCRIBE_TIMEOUT_SECONDS, TimeUnit.SECONDS)) + assertThat(itemCountQueryIssued.await(SUBSCRIBE_TIMEOUT_SECONDS, TimeUnit.SECONDS)) .as("the item-count query should be issued while the unrelated summaries are still pending") .isTrue(); } finally { From 5528323b1cbb7cdbd52ba39107190aba0bfdb0a9 Mon Sep 17 00:00:00 2001 From: Daniel Dimenshtein Date: Fri, 4 Sep 2026 12:25:33 +0300 Subject: [PATCH 6/6] test(datasets): wait via Awaitility instead of hand-rolled latch coordination Hand-rolled CountDownLatch waiting had to get the deadline and interrupt handling right itself, and a missed expectation left a worker parked on a bare await() -- the failure mode that turns into a hung CI job. Awaitility now owns every wait in the non-blocking test: untilTrue for the assertion, and the same for the gate holding the unrelated summaries, so awaitQuietly is gone. A timeout now reports the named condition instead of a bare assertion failure. Also import AtomicBoolean rather than spelling it fully qualified. Co-Authored-By: Claude Opus 5 (1M context) --- .../domain/DatasetServiceEnrichmentTest.java | 50 ++++++++++--------- 1 file changed, 27 insertions(+), 23 deletions(-) diff --git a/apps/opik-backend/src/test/java/com/comet/opik/domain/DatasetServiceEnrichmentTest.java b/apps/opik-backend/src/test/java/com/comet/opik/domain/DatasetServiceEnrichmentTest.java index 60d15bf98e5..fd9e6681b75 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/domain/DatasetServiceEnrichmentTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/domain/DatasetServiceEnrichmentTest.java @@ -11,6 +11,7 @@ import com.comet.opik.infrastructure.auth.RequestContext; import com.google.common.eventbus.EventBus; import jakarta.inject.Provider; +import org.awaitility.Awaitility; import org.jdbi.v3.core.Handle; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.DisplayName; @@ -35,6 +36,7 @@ import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.function.Function; import static org.assertj.core.api.Assertions.assertThat; @@ -605,8 +607,8 @@ void enrichmentDefersItemCountUntilVersionsAreKnown() { stubPage(datasetId); stubVersioningEnabled(); - var versionLookupDone = new java.util.concurrent.atomic.AtomicBoolean(false); - var itemCountSawResolvedVersions = new java.util.concurrent.atomic.AtomicBoolean(false); + var versionLookupDone = new AtomicBoolean(false); + var itemCountSawResolvedVersions = new AtomicBoolean(false); when(datasetVersionDAO.findLatestVersionsByDatasetIds(anySet(), any())).thenAnswer(invocation -> { versionLookupDone.set(true); @@ -634,27 +636,30 @@ void enrichmentDoesNotBlockItemCountOnUnrelatedSummaries() throws Exception { stubLatestVersions(); // The item count depends only on the versions, so the query must be issued while the two unrelated - // summaries are still in flight. The latch counts down in the Mockito answer -- i.e. when the DAO - // method is invoked, which is the point the narrowed id set is computed and handed over, not the + // summaries are still in flight. itemCountQueryIssued is set in the Mockito answer -- i.e. when the + // DAO method is invoked, which is the point the narrowed id set is computed and handed over, not the // later subscription to the returned Flux. That is the property under test: if the count were - // chained off the whole zip it would never be requested and the latch would time out. - var itemCountQueryIssued = new CountDownLatch(1); - var releaseSummaries = new CountDownLatch(1); + // chained off the whole zip it would never be requested while the summaries are held. + var itemCountQueryIssued = new AtomicBoolean(false); + var summariesReleased = new AtomicBoolean(false); + // Both unrelated summaries park until the assertion below has run. Awaitility owns every wait in this + // test -- it enforces the deadline and surfaces a ConditionTimeoutException naming the unmet condition, + // rather than leaving a worker blocked on a bare await() if the expectation never holds. Flux heldExperiments = Flux.defer(() -> { - awaitQuietly(releaseSummaries); + awaitReleased(summariesReleased); return Flux.empty(); }).subscribeOn(Schedulers.boundedElastic()); Flux heldOptimizations = Flux.defer(() -> { - awaitQuietly(releaseSummaries); + awaitReleased(summariesReleased); return Flux.empty(); }).subscribeOn(Schedulers.boundedElastic()); when(experimentItemDAO.findExperimentSummaryByDatasetIds(anySet())).thenReturn(heldExperiments); when(optimizationDAO.findOptimizationSummaryByDatasetIds(anySet())).thenReturn(heldOptimizations); when(datasetItemDAO.findDatasetItemSummaryByDatasetIds(anySet())).thenAnswer(invocation -> { - itemCountQueryIssued.countDown(); + itemCountQueryIssued.set(true); return Flux.just(new DatasetItemSummary(datasetId, 6)); }); @@ -662,26 +667,25 @@ void enrichmentDoesNotBlockItemCountOnUnrelatedSummaries() throws Exception { .supplyAsync(() -> service.find(1, 10, DatasetCriteria.builder().build(), List.of())); try { - assertThat(itemCountQueryIssued.await(SUBSCRIBE_TIMEOUT_SECONDS, TimeUnit.SECONDS)) - .as("the item-count query should be issued while the unrelated summaries are still pending") - .isTrue(); + Awaitility.await("the item-count query is issued while the unrelated summaries are still pending") + .atMost(SUBSCRIBE_TIMEOUT_SECONDS, TimeUnit.SECONDS) + .pollInterval(50, TimeUnit.MILLISECONDS) + .untilTrue(itemCountQueryIssued); } finally { - releaseSummaries.countDown(); + // Always release, so a failed assertion reports the real cause instead of being masked by the + // held summaries timing out. + summariesReleased.set(true); } var actual = enrichment.get(WORKER_RELEASE_TIMEOUT_SECONDS, TimeUnit.SECONDS); assertThat(itemsCountById(actual.content())).containsExactlyInAnyOrderEntriesOf(Map.of(datasetId, 6L)); } - private static void awaitQuietly(CountDownLatch latch) { - try { - if (!latch.await(WORKER_RELEASE_TIMEOUT_SECONDS, TimeUnit.SECONDS)) { - throw new IllegalStateException("summaries were never released"); - } - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - throw new IllegalStateException(e); - } + private static void awaitReleased(AtomicBoolean released) { + Awaitility.await("the held summaries are released") + .atMost(WORKER_RELEASE_TIMEOUT_SECONDS, TimeUnit.SECONDS) + .pollInterval(50, TimeUnit.MILLISECONDS) + .untilTrue(released); } private void stubVersioningEnabled() {