From 97c1bf01ff511c4db74dc8a81045447b009bec29 Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Wed, 7 Aug 2024 09:39:07 -0700 Subject: [PATCH 01/40] QueryGroup Resource Tracking framework and implementation (#13897) * initial code for the sandbox resource tracking and cancellation framework Signed-off-by: Kiran Prakash * Fix Failing Tests Signed-off-by: Kiran Prakash * spotless Apply Signed-off-by: Kiran Prakash * Update SandboxService.java Signed-off-by: Kiran Prakash * Update SandboxService.java Signed-off-by: Kiran Prakash * Update SandboxTask.java Signed-off-by: Kiran Prakash * Add java docs Signed-off-by: Kiran Prakash * spotless Signed-off-by: Kiran Prakash * javadocs Signed-off-by: Kiran Prakash * javadocs Signed-off-by: Kiran Prakash * java docs Signed-off-by: Kiran Prakash * Update AbstractTaskCancellation.java Signed-off-by: Kiran Prakash * Update SandboxModule.java Signed-off-by: Kiran Prakash * Some tests and stubs Signed-off-by: Kiran Prakash * spotless Signed-off-by: Kiran Prakash * :server:testingConventions Signed-off-by: Kiran Prakash * Update AbstractTaskCancellation.java Signed-off-by: Kiran Prakash * more tests Signed-off-by: Kiran Prakash * addressing comments Signed-off-by: Kiran Prakash * revert some accidentally pushed files Signed-off-by: Kiran Prakash * resolve flakiness Signed-off-by: Kiran Prakash * renaming sandbox to querygroup and adjusting code based on merged PRs Signed-off-by: Kiran Prakash * jvm to memory Signed-off-by: Kiran Prakash * missing java docs Signed-off-by: Kiran Prakash * spotless Signed-off-by: Kiran Prakash * Update CHANGELOG.md Signed-off-by: Kiran Prakash * pluck cancellation changes out of this PR Signed-off-by: Kiran Prakash * remove unused Signed-off-by: Kiran Prakash * remove cancellation related code and add more tests coverage Signed-off-by: Kiran Prakash * us only memory and not jvm Signed-off-by: Kiran Prakash * test conventions Signed-off-by: Kiran Prakash * Bring back enum Signed-off-by: Kiran Prakash * Update SearchBackpressureService.java Signed-off-by: Kiran Prakash * revert changes Signed-off-by: Kiran Prakash * revert changes Signed-off-by: Kiran Prakash * all required changes Signed-off-by: Kiran Prakash * Update CHANGELOG.md Signed-off-by: Kiran Prakash * cleanups Signed-off-by: Kiran Prakash * Delete QueryGroupService.java Signed-off-by: Kiran Prakash * cleanups Signed-off-by: Kiran Prakash * Update QueryGroupLevelResourceUsageViewTests.java Signed-off-by: Kiran Prakash * Update QueryGroupLevelResourceUsageViewTests.java Signed-off-by: Kiran Prakash * Update QueryGroupResourceUsageTrackerService.java Signed-off-by: Kiran Prakash * Update QueryGroupResourceUsageTrackerService.java Signed-off-by: Kiran Prakash * Update QueryGroupResourceUsageTrackerService.java Signed-off-by: Kiran Prakash * Update CHANGELOG.md Signed-off-by: Kiran Prakash * rebasing with latest main Signed-off-by: Kiran Prakash * remove experimental Signed-off-by: Kiran Prakash * remove queryGroupId Signed-off-by: Kiran Prakash * Update QueryGroupResourceUsageTrackerService.java Signed-off-by: Kiran Prakash * change code comments Signed-off-by: Kiran Prakash * remmove QueryGroupUsageTracker Signed-off-by: Kiran Prakash * Update QueryGroupResourceUsageTrackerService.java Signed-off-by: Kiran Prakash * Update QueryGroupResourceUsageTrackerService.java Signed-off-by: Kiran Prakash * remove QueryGroupTestHelpers Signed-off-by: Kiran Prakash * cleanups Signed-off-by: Kiran Prakash * remove queryGroupHelper Signed-off-by: Kiran Prakash * Update ResourceTypeTests.java Signed-off-by: Kiran Prakash * extend OpenSearchTestCase Signed-off-by: Kiran Prakash * pr comments Signed-off-by: Kiran Prakash * Update CHANGELOG.md Signed-off-by: Kiran Prakash * Update QueryGroupResourceUsageTrackerServiceTests.java Signed-off-by: Kiran Prakash * Update ResourceTypeTests.java Signed-off-by: Kiran Prakash * Update ResourceTypeTests.java Signed-off-by: Kiran Prakash * Update ResourceType.java Signed-off-by: Kiran Prakash * Update ResourceType.java Signed-off-by: Kiran Prakash --------- Signed-off-by: Kiran Prakash --- CHANGELOG.md | 1 + .../opensearch/cluster/metadata/Metadata.java | 2 +- .../org/opensearch/search/ResourceType.java | 21 ++- .../wlm/QueryGroupLevelResourceUsageView.java | 50 +++++++ ...QueryGroupResourceUsageTrackerService.java | 84 ++++++++++++ .../opensearch/wlm/tracker/package-info.java | 12 ++ .../opensearch/search/ResourceTypeTests.java | 52 ++++++++ ...QueryGroupLevelResourceUsageViewTests.java | 64 +++++++++ ...GroupResourceUsageTrackerServiceTests.java | 126 ++++++++++++++++++ 9 files changed, 408 insertions(+), 4 deletions(-) create mode 100644 server/src/main/java/org/opensearch/wlm/QueryGroupLevelResourceUsageView.java create mode 100644 server/src/main/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerService.java create mode 100644 server/src/main/java/org/opensearch/wlm/tracker/package-info.java create mode 100644 server/src/test/java/org/opensearch/search/ResourceTypeTests.java create mode 100644 server/src/test/java/org/opensearch/wlm/QueryGroupLevelResourceUsageViewTests.java create mode 100644 server/src/test/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerServiceTests.java diff --git a/CHANGELOG.md b/CHANGELOG.md index be5e5598b09c2..3e83a5bf9b4cb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,6 +14,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), - Add ThreadContextPermission for stashAndMergeHeaders and stashWithOrigin ([#15039](https://github.com/opensearch-project/OpenSearch/pull/15039)) - [Concurrent Segment Search] Support composite aggregations with scripting ([#15072](https://github.com/opensearch-project/OpenSearch/pull/15072)) - Add `rangeQuery` and `regexpQuery` for `constant_keyword` field type ([#14711](https://github.com/opensearch-project/OpenSearch/pull/14711)) +- [Workload Management] QueryGroup resource tracking framework changes ([#13897](https://github.com/opensearch-project/OpenSearch/pull/13897)) ### Dependencies - Bump `netty` from 4.1.111.Final to 4.1.112.Final ([#15081](https://github.com/opensearch-project/OpenSearch/pull/15081)) diff --git a/server/src/main/java/org/opensearch/cluster/metadata/Metadata.java b/server/src/main/java/org/opensearch/cluster/metadata/Metadata.java index 440b9e267cf0a..09bef2ddf9ee6 100644 --- a/server/src/main/java/org/opensearch/cluster/metadata/Metadata.java +++ b/server/src/main/java/org/opensearch/cluster/metadata/Metadata.java @@ -1391,7 +1391,7 @@ public Builder put(final QueryGroup queryGroup) { return queryGroups(existing); } - private Map getQueryGroups() { + public Map getQueryGroups() { return Optional.ofNullable(this.customs.get(QueryGroupMetadata.TYPE)) .map(o -> (QueryGroupMetadata) o) .map(QueryGroupMetadata::queryGroups) diff --git a/server/src/main/java/org/opensearch/search/ResourceType.java b/server/src/main/java/org/opensearch/search/ResourceType.java index fe5ce4dd2bb50..0cba2222a6e20 100644 --- a/server/src/main/java/org/opensearch/search/ResourceType.java +++ b/server/src/main/java/org/opensearch/search/ResourceType.java @@ -10,21 +10,26 @@ import org.opensearch.common.annotation.PublicApi; import org.opensearch.core.common.io.stream.StreamOutput; +import org.opensearch.core.tasks.resourcetracker.ResourceStats; +import org.opensearch.tasks.Task; import java.io.IOException; +import java.util.function.Function; /** * Enum to hold the resource type */ @PublicApi(since = "2.x") public enum ResourceType { - CPU("cpu"), - MEMORY("memory"); + CPU("cpu", task -> task.getTotalResourceUtilization(ResourceStats.CPU)), + MEMORY("memory", task -> task.getTotalResourceUtilization(ResourceStats.MEMORY)); private final String name; + private final Function getResourceUsage; - ResourceType(String name) { + ResourceType(String name, Function getResourceUsage) { this.name = name; + this.getResourceUsage = getResourceUsage; } /** @@ -48,4 +53,14 @@ public static void writeTo(StreamOutput out, ResourceType resourceType) throws I public String getName() { return name; } + + /** + * Gets the resource usage for a given resource type and task. + * + * @param task the task for which to calculate resource usage + * @return the resource usage + */ + public long getResourceUsage(Task task) { + return getResourceUsage.apply(task); + } } diff --git a/server/src/main/java/org/opensearch/wlm/QueryGroupLevelResourceUsageView.java b/server/src/main/java/org/opensearch/wlm/QueryGroupLevelResourceUsageView.java new file mode 100644 index 0000000000000..2fd743dc3f83f --- /dev/null +++ b/server/src/main/java/org/opensearch/wlm/QueryGroupLevelResourceUsageView.java @@ -0,0 +1,50 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.wlm; + +import org.opensearch.search.ResourceType; +import org.opensearch.tasks.Task; + +import java.util.List; +import java.util.Map; + +/** + * Represents the point in time view of resource usage of a QueryGroup and + * has a 1:1 relation with a QueryGroup. + * This class holds the resource usage data and the list of active tasks. + */ +public class QueryGroupLevelResourceUsageView { + // resourceUsage holds the resource usage data for a QueryGroup at a point in time + private final Map resourceUsage; + // activeTasks holds the list of active tasks for a QueryGroup at a point in time + private final List activeTasks; + + public QueryGroupLevelResourceUsageView(Map resourceUsage, List activeTasks) { + this.resourceUsage = resourceUsage; + this.activeTasks = activeTasks; + } + + /** + * Returns the resource usage data. + * + * @return The map of resource usage data + */ + public Map getResourceUsageData() { + return resourceUsage; + } + + /** + * Returns the list of active tasks. + * + * @return The list of active tasks + */ + public List getActiveTasks() { + return activeTasks; + } +} diff --git a/server/src/main/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerService.java b/server/src/main/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerService.java new file mode 100644 index 0000000000000..bfbf5d8a452d1 --- /dev/null +++ b/server/src/main/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerService.java @@ -0,0 +1,84 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.wlm.tracker; + +import org.opensearch.search.ResourceType; +import org.opensearch.tasks.Task; +import org.opensearch.tasks.TaskResourceTrackingService; +import org.opensearch.wlm.QueryGroupLevelResourceUsageView; +import org.opensearch.wlm.QueryGroupTask; + +import java.util.EnumMap; +import java.util.EnumSet; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; + +/** + * This class tracks resource usage per QueryGroup + */ +public class QueryGroupResourceUsageTrackerService { + + public static final EnumSet TRACKED_RESOURCES = EnumSet.allOf(ResourceType.class); + private final TaskResourceTrackingService taskResourceTrackingService; + + /** + * QueryGroupResourceTrackerService constructor + * + * @param taskResourceTrackingService Service that helps track resource usage of tasks running on a node. + */ + public QueryGroupResourceUsageTrackerService(TaskResourceTrackingService taskResourceTrackingService) { + this.taskResourceTrackingService = taskResourceTrackingService; + } + + /** + * Constructs a map of QueryGroupLevelResourceUsageView instances for each QueryGroup. + * + * @return Map of QueryGroup views + */ + public Map constructQueryGroupLevelUsageViews() { + final Map> tasksByQueryGroup = getTasksGroupedByQueryGroup(); + final Map queryGroupViews = new HashMap<>(); + + // Iterate over each QueryGroup entry + for (Map.Entry> queryGroupEntry : tasksByQueryGroup.entrySet()) { + // Compute the QueryGroup usage + final EnumMap queryGroupUsage = new EnumMap<>(ResourceType.class); + for (ResourceType resourceType : TRACKED_RESOURCES) { + long queryGroupResourceUsage = 0; + for (Task task : queryGroupEntry.getValue()) { + queryGroupResourceUsage += resourceType.getResourceUsage(task); + } + queryGroupUsage.put(resourceType, queryGroupResourceUsage); + } + + // Add to the QueryGroup View + queryGroupViews.put( + queryGroupEntry.getKey(), + new QueryGroupLevelResourceUsageView(queryGroupUsage, queryGroupEntry.getValue()) + ); + } + return queryGroupViews; + } + + /** + * Groups tasks by their associated QueryGroup. + * + * @return Map of tasks grouped by QueryGroup + */ + private Map> getTasksGroupedByQueryGroup() { + return taskResourceTrackingService.getResourceAwareTasks() + .values() + .stream() + .filter(QueryGroupTask.class::isInstance) + .map(QueryGroupTask.class::cast) + .collect(Collectors.groupingBy(QueryGroupTask::getQueryGroupId, Collectors.mapping(task -> (Task) task, Collectors.toList()))); + } +} diff --git a/server/src/main/java/org/opensearch/wlm/tracker/package-info.java b/server/src/main/java/org/opensearch/wlm/tracker/package-info.java new file mode 100644 index 0000000000000..86efc99355d3d --- /dev/null +++ b/server/src/main/java/org/opensearch/wlm/tracker/package-info.java @@ -0,0 +1,12 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +/** + * QueryGroup resource tracking artifacts + */ +package org.opensearch.wlm.tracker; diff --git a/server/src/test/java/org/opensearch/search/ResourceTypeTests.java b/server/src/test/java/org/opensearch/search/ResourceTypeTests.java new file mode 100644 index 0000000000000..78827b8b1bdad --- /dev/null +++ b/server/src/test/java/org/opensearch/search/ResourceTypeTests.java @@ -0,0 +1,52 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.search; + +import org.opensearch.action.search.SearchShardTask; +import org.opensearch.core.tasks.resourcetracker.ResourceStats; +import org.opensearch.tasks.CancellableTask; +import org.opensearch.test.OpenSearchTestCase; + +import static org.junit.Assert.assertThrows; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +public class ResourceTypeTests extends OpenSearchTestCase { + + public void testFromName() { + assertSame(ResourceType.CPU, ResourceType.fromName("cpu")); + assertThrows(IllegalArgumentException.class, () -> { ResourceType.fromName("CPU"); }); + assertThrows(IllegalArgumentException.class, () -> { ResourceType.fromName("Cpu"); }); + + assertSame(ResourceType.MEMORY, ResourceType.fromName("memory")); + assertThrows(IllegalArgumentException.class, () -> { ResourceType.fromName("Memory"); }); + assertThrows(IllegalArgumentException.class, () -> { ResourceType.fromName("MEMORY"); }); + assertThrows(IllegalArgumentException.class, () -> { ResourceType.fromName("JVM"); }); + assertThrows(IllegalArgumentException.class, () -> { ResourceType.fromName("Heap"); }); + assertThrows(IllegalArgumentException.class, () -> { ResourceType.fromName("Disk"); }); + } + + public void testGetName() { + assertEquals("cpu", ResourceType.CPU.getName()); + assertEquals("memory", ResourceType.MEMORY.getName()); + } + + public void testGetResourceUsage() { + SearchShardTask mockTask = createMockTask(SearchShardTask.class, 100, 200); + assertEquals(100, ResourceType.CPU.getResourceUsage(mockTask)); + assertEquals(200, ResourceType.MEMORY.getResourceUsage(mockTask)); + } + + private T createMockTask(Class type, long cpuUsage, long heapUsage) { + T task = mock(type); + when(task.getTotalResourceUtilization(ResourceStats.CPU)).thenReturn(cpuUsage); + when(task.getTotalResourceUtilization(ResourceStats.MEMORY)).thenReturn(heapUsage); + return task; + } +} diff --git a/server/src/test/java/org/opensearch/wlm/QueryGroupLevelResourceUsageViewTests.java b/server/src/test/java/org/opensearch/wlm/QueryGroupLevelResourceUsageViewTests.java new file mode 100644 index 0000000000000..7f6419505fec2 --- /dev/null +++ b/server/src/test/java/org/opensearch/wlm/QueryGroupLevelResourceUsageViewTests.java @@ -0,0 +1,64 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.wlm; + +import org.opensearch.action.search.SearchAction; +import org.opensearch.core.tasks.TaskId; +import org.opensearch.search.ResourceType; +import org.opensearch.tasks.Task; +import org.opensearch.test.OpenSearchTestCase; + +import java.util.Collections; +import java.util.List; +import java.util.Map; + +public class QueryGroupLevelResourceUsageViewTests extends OpenSearchTestCase { + Map resourceUsage; + List activeTasks; + + public void setUp() throws Exception { + super.setUp(); + resourceUsage = Map.of(ResourceType.fromName("memory"), 34L, ResourceType.fromName("cpu"), 12L); + activeTasks = List.of(getRandomTask(4321)); + } + + public void testGetResourceUsageData() { + QueryGroupLevelResourceUsageView queryGroupLevelResourceUsageView = new QueryGroupLevelResourceUsageView( + resourceUsage, + activeTasks + ); + Map resourceUsageData = queryGroupLevelResourceUsageView.getResourceUsageData(); + assertTrue(assertResourceUsageData(resourceUsageData)); + } + + public void testGetActiveTasks() { + QueryGroupLevelResourceUsageView queryGroupLevelResourceUsageView = new QueryGroupLevelResourceUsageView( + resourceUsage, + activeTasks + ); + List activeTasks = queryGroupLevelResourceUsageView.getActiveTasks(); + assertEquals(1, activeTasks.size()); + assertEquals(4321, activeTasks.get(0).getId()); + } + + private boolean assertResourceUsageData(Map resourceUsageData) { + return resourceUsageData.get(ResourceType.fromName("memory")) == 34L && resourceUsageData.get(ResourceType.fromName("cpu")) == 12L; + } + + private Task getRandomTask(long id) { + return new Task( + id, + "transport", + SearchAction.NAME, + "test description", + new TaskId(randomLong() + ":" + randomLong()), + Collections.emptyMap() + ); + } +} diff --git a/server/src/test/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerServiceTests.java b/server/src/test/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerServiceTests.java new file mode 100644 index 0000000000000..967119583c25f --- /dev/null +++ b/server/src/test/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerServiceTests.java @@ -0,0 +1,126 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.wlm.tracker; + +import org.opensearch.action.search.SearchShardTask; +import org.opensearch.action.search.SearchTask; +import org.opensearch.common.util.concurrent.ThreadContext; +import org.opensearch.core.tasks.resourcetracker.ResourceStats; +import org.opensearch.search.ResourceType; +import org.opensearch.tasks.CancellableTask; +import org.opensearch.tasks.Task; +import org.opensearch.tasks.TaskResourceTrackingService; +import org.opensearch.test.OpenSearchTestCase; +import org.opensearch.threadpool.TestThreadPool; +import org.opensearch.threadpool.ThreadPool; +import org.opensearch.wlm.QueryGroupLevelResourceUsageView; +import org.opensearch.wlm.QueryGroupTask; +import org.junit.After; +import org.junit.Before; + +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; + +import static org.opensearch.wlm.QueryGroupTask.QUERY_GROUP_ID_HEADER; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +public class QueryGroupResourceUsageTrackerServiceTests extends OpenSearchTestCase { + TestThreadPool threadPool; + TaskResourceTrackingService mockTaskResourceTrackingService; + QueryGroupResourceUsageTrackerService queryGroupResourceUsageTrackerService; + + @Before + public void setup() { + threadPool = new TestThreadPool(getTestName()); + mockTaskResourceTrackingService = mock(TaskResourceTrackingService.class); + queryGroupResourceUsageTrackerService = new QueryGroupResourceUsageTrackerService(mockTaskResourceTrackingService); + } + + @After + public void cleanup() { + ThreadPool.terminate(threadPool, 5, TimeUnit.SECONDS); + } + + public void testConstructQueryGroupLevelViews_CreatesQueryGroupLevelUsageView_WhenTasksArePresent() { + List queryGroupIds = List.of("queryGroup1", "queryGroup2", "queryGroup3"); + + Map activeSearchShardTasks = createActiveSearchShardTasks(queryGroupIds); + when(mockTaskResourceTrackingService.getResourceAwareTasks()).thenReturn(activeSearchShardTasks); + Map stringQueryGroupLevelResourceUsageViewMap = queryGroupResourceUsageTrackerService + .constructQueryGroupLevelUsageViews(); + + for (String queryGroupId : queryGroupIds) { + assertEquals( + 400, + (long) stringQueryGroupLevelResourceUsageViewMap.get(queryGroupId).getResourceUsageData().get(ResourceType.MEMORY) + ); + assertEquals(2, stringQueryGroupLevelResourceUsageViewMap.get(queryGroupId).getActiveTasks().size()); + } + } + + public void testConstructQueryGroupLevelViews_CreatesQueryGroupLevelUsageView_WhenTasksAreNotPresent() { + Map stringQueryGroupLevelResourceUsageViewMap = queryGroupResourceUsageTrackerService + .constructQueryGroupLevelUsageViews(); + assertTrue(stringQueryGroupLevelResourceUsageViewMap.isEmpty()); + } + + public void testConstructQueryGroupLevelUsageViews_WithTasksHavingDifferentResourceUsage() { + Map activeSearchShardTasks = new HashMap<>(); + activeSearchShardTasks.put(1L, createMockTask(SearchShardTask.class, 100, 200, "queryGroup1")); + activeSearchShardTasks.put(2L, createMockTask(SearchShardTask.class, 200, 400, "queryGroup1")); + when(mockTaskResourceTrackingService.getResourceAwareTasks()).thenReturn(activeSearchShardTasks); + + Map queryGroupViews = queryGroupResourceUsageTrackerService + .constructQueryGroupLevelUsageViews(); + + assertEquals(600, (long) queryGroupViews.get("queryGroup1").getResourceUsageData().get(ResourceType.MEMORY)); + assertEquals(2, queryGroupViews.get("queryGroup1").getActiveTasks().size()); + } + + private Map createActiveSearchShardTasks(List queryGroupIds) { + Map activeSearchShardTasks = new HashMap<>(); + long task_id = 0; + for (String queryGroupId : queryGroupIds) { + for (int i = 0; i < 2; i++) { + activeSearchShardTasks.put(++task_id, createMockTask(SearchShardTask.class, 100, 200, queryGroupId)); + } + } + return activeSearchShardTasks; + } + + private T createMockTask(Class type, long cpuUsage, long heapUsage, String queryGroupId) { + T task = mock(type); + if (task instanceof SearchTask || task instanceof SearchShardTask) { + // Stash the current thread context to ensure that any existing context is preserved and restored after setting the query group + // ID. + try (ThreadContext.StoredContext ignore = threadPool.getThreadContext().stashContext()) { + threadPool.getThreadContext().putHeader(QUERY_GROUP_ID_HEADER, queryGroupId); + ((QueryGroupTask) task).setQueryGroupId(threadPool.getThreadContext()); + } + } + when(task.getTotalResourceUtilization(ResourceStats.CPU)).thenReturn(cpuUsage); + when(task.getTotalResourceUtilization(ResourceStats.MEMORY)).thenReturn(heapUsage); + when(task.getStartTimeNanos()).thenReturn((long) 0); + + AtomicBoolean isCancelled = new AtomicBoolean(false); + doAnswer(invocation -> { + isCancelled.set(true); + return null; + }).when(task).cancel(anyString()); + doAnswer(invocation -> isCancelled.get()).when(task).isCancelled(); + + return task; + } +} From 348c04e7a32e13ea040a1d2e0459c03da9ec0c2c Mon Sep 17 00:00:00 2001 From: Jay Deng Date: Wed, 7 Aug 2024 11:40:52 -0700 Subject: [PATCH 02/40] Fix CHANGELOG for #15054 (#15150) Signed-off-by: Jay Deng --- CHANGELOG-3.0.md | 1 - CHANGELOG.md | 1 + 2 files changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG-3.0.md b/CHANGELOG-3.0.md index 78e93eed0158a..48d978bede420 100644 --- a/CHANGELOG-3.0.md +++ b/CHANGELOG-3.0.md @@ -13,7 +13,6 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), - GHA to verify checklist items completion in PR descriptions ([#10800](https://github.com/opensearch-project/OpenSearch/pull/10800)) - Allow to pass the list settings through environment variables (like [], ["a", "b", "c"], ...) ([#10625](https://github.com/opensearch-project/OpenSearch/pull/10625)) - Views, simplify data access and manipulation by providing a virtual layer over one or more indices ([#11957](https://github.com/opensearch-project/OpenSearch/pull/11957)) -- Add took time to request nodes stats ([#15054](https://github.com/opensearch-project/OpenSearch/pull/15054)) ### Dependencies diff --git a/CHANGELOG.md b/CHANGELOG.md index 3e83a5bf9b4cb..f44949bf38511 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,6 +14,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), - Add ThreadContextPermission for stashAndMergeHeaders and stashWithOrigin ([#15039](https://github.com/opensearch-project/OpenSearch/pull/15039)) - [Concurrent Segment Search] Support composite aggregations with scripting ([#15072](https://github.com/opensearch-project/OpenSearch/pull/15072)) - Add `rangeQuery` and `regexpQuery` for `constant_keyword` field type ([#14711](https://github.com/opensearch-project/OpenSearch/pull/14711)) +- Add took time to request nodes stats ([#15054](https://github.com/opensearch-project/OpenSearch/pull/15054)) - [Workload Management] QueryGroup resource tracking framework changes ([#13897](https://github.com/opensearch-project/OpenSearch/pull/13897)) ### Dependencies From 7f72a6e6580e42740a938c35dc00a6f88e6089df Mon Sep 17 00:00:00 2001 From: Andrew Ross Date: Wed, 7 Aug 2024 15:48:31 -0500 Subject: [PATCH 03/40] Add 'ShardManagement:*' labels to Cluster Manager triage search (#14234) Signed-off-by: Andrew Ross --- TRIAGING.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/TRIAGING.md b/TRIAGING.md index 53ef77de49159..dddcbc15394ab 100644 --- a/TRIAGING.md +++ b/TRIAGING.md @@ -31,8 +31,8 @@ Meeting structure may vary slightly, but the general structure is as follows: - [Search](https://github.com/opensearch-project/OpenSearch/issues?q=is%3Aissue+is%3Aopen+label%3Auntriaged+label%3A%22Search%22%2C%22Search%3ARemote+Search%22%2C%22Search%3AResiliency%22%2C%22Search%3APerformance%22%2C%22Search%3ARelevance%22%2C%22Search%3AAggregations%22%2C%22Search%3AQuery+Capabilities%22%2C%22Search%3AQuery+Insights%22%2C%22Search%3ASearchable+Snapshots%22%2C%22Search%3AUser+Behavior+Insights%22) - [Indexing](https://github.com/opensearch-project/OpenSearch/issues?q=is%3Aissue+is%3Aopen+label%3Auntriaged+label%3A%22Indexing%3AReplication%22%2C%22Indexing%22%2C%22Indexing%3APerformance%22%2C%22Indexing+%26+Search%22%2C) - [Storage](https://github.com/opensearch-project/OpenSearch/issues?q=is%3Aissue+is%3Aopen+label%3Auntriaged+label%3AStorage%2C%22Storage%3AResiliency%22%2C%22Storage%3APerformance%22%2C%22Storage%3ASnapshots%22%2C%22Storage%3ARemote%22%2C%22Storage%3ADurability%22) - - [Cluster Manager](https://github.com/opensearch-project/OpenSearch/issues?q=is%3Aissue+is%3Aopen+label%3Auntriaged+label%3A%22Cluster+Manager%22%2C%22ClusterManager%3ARemoteState%22) - - [Core](https://github.com/opensearch-project/OpenSearch/issues?q=is%3Aissue+is%3Aopen+label%3Auntriaged+-label%3A%22Search%22%2C%22Search%3ARemote+Search%22%2C%22Search%3AResiliency%22%2C%22Search%3APerformance%22%2C%22Search%3ARelevance%22%2C%22Search%3AAggregations%22%2C%22Search%3AQuery+Capabilities%22%2C%22Search%3AQuery+Insights%22%2C%22Search%3ASearchable+Snapshots%22%2C%22Search%3AUser+Behavior+Insights%22%2C%22Storage%22%2C%22Storage%3AResiliency%22%2C%22Storage%3APerformance%22%2C%22Storage%3ASnapshots%22%2C%22Storage%3ARemote%22%2C%22Storage%3ADurability%22%2C%22Cluster+Manager%22%2C%22ClusterManager%3ARemoteState%22%2C%22Indexing%3AReplication%22%2C%22Indexing%22%2C%22Indexing%3APerformance%22%2C%22Indexing+%26+Search%22) + - [Cluster Manager](https://github.com/opensearch-project/OpenSearch/issues?q=is%3Aissue+is%3Aopen+label%3Auntriaged+label%3A%22Cluster+Manager%22%2C%22ClusterManager%3ARemoteState%22%2C%22ShardManagement%3AResiliency%22%2C%22ShardManagement%3AInsights%22%2C%22ShardManagement%3ASizing%22%2C%22ShardManagement%3APerformance%22%2C%22ShardManagement%3APlacement%22%2C%22ShardManagement%3ARouting%22) + - [Core](https://github.com/opensearch-project/OpenSearch/issues?q=is%3Aissue+is%3Aopen+label%3Auntriaged+-label%3A%22Search%22%2C%22Search%3ARemote+Search%22%2C%22Search%3AResiliency%22%2C%22Search%3APerformance%22%2C%22Search%3ARelevance%22%2C%22Search%3AAggregations%22%2C%22Search%3AQuery+Capabilities%22%2C%22Search%3AQuery+Insights%22%2C%22Search%3ASearchable+Snapshots%22%2C%22Search%3AUser+Behavior+Insights%22%2C%22Storage%22%2C%22Storage%3AResiliency%22%2C%22Storage%3APerformance%22%2C%22Storage%3ASnapshots%22%2C%22Storage%3ARemote%22%2C%22Storage%3ADurability%22%2C%22Cluster+Manager%22%2C%22ClusterManager%3ARemoteState%22%2C%22ShardManagement%3AResiliency%22%2C%22ShardManagement%3AInsights%22%2C%22ShardManagement%3ASizing%22%2C%22ShardManagement%3APerformance%22%2C%22ShardManagement%3APlacement%22%2C%22ShardManagement%3ARouting%22%2C%22Indexing%3AReplication%22%2C%22Indexing%22%2C%22Indexing%3APerformance%22%2C%22Indexing+%26+Search%22) 5. **Attendee Requests:** An opportunity for any meeting member to request consideration of an issue or pull request. 6. **Open Discussion:** Attendees can bring up any topics not already covered by filed issues or pull requests. 7. **Review of Old Untriaged Issues:** Look at all [untriaged issues older than 14 days](https://peternied.github.io/redirect/issue_search.html?owner=opensearch-project&repo=OpenSearch&tag=untriaged&created-since-days=14) to prevent issues from falling through the cracks. From 26ff6353669a1560c61909046a756cc9a9aa1ef5 Mon Sep 17 00:00:00 2001 From: Rishabh Singh Date: Wed, 7 Aug 2024 16:55:03 -0700 Subject: [PATCH 04/40] fix missing curly brace (#15157) Signed-off-by: Rishabh Singh --- .github/workflows/benchmark-pull-request.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/benchmark-pull-request.yml b/.github/workflows/benchmark-pull-request.yml index 1096014e4a291..0173b7e35c64d 100644 --- a/.github/workflows/benchmark-pull-request.yml +++ b/.github/workflows/benchmark-pull-request.yml @@ -62,6 +62,7 @@ jobs: } if (benchmarkConfigs[configId].hasOwnProperty('baseline_cluster_config')) { core.exportVariable('BASELINE_CLUSTER_CONFIG', benchmarkConfigs[configId]['baseline_cluster_config']); + } - name: Post invalid format comment if: steps.check_comment.outputs.invalid == 'true' uses: actions/github-script@v7 From c6189a931d437ffe3fc8eb3784f8adfc2ddcdff9 Mon Sep 17 00:00:00 2001 From: Gaurav Bafna <85113518+gbbafna@users.noreply.github.com> Date: Thu, 8 Aug 2024 13:46:22 +0530 Subject: [PATCH 05/40] Add Varun Bansal as maintainer (#15163) Signed-off-by: Gaurav Bafna --- .github/CODEOWNERS | 2 +- MAINTAINERS.md | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index fb7d73f599670..e75e22f30431d 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -11,7 +11,7 @@ # 3. Use the command palette to run the CODEOWNERS: Show owners of current file command, which will display all code owners for the current file. # Default ownership for all repo files -* @anasalkouz @andrross @ashking94 @Bukhtawar @CEHENKLE @dblock @dbwiddis @gbbafna @kotwanikunal @mch2 @msfroh @nknize @owaiskazi19 @reta @Rishikesh1159 @sachinpkale @saratvemulapalli @shwetathareja @sohami @VachaShah +* @anasalkouz @andrross @ashking94 @Bukhtawar @CEHENKLE @dblock @dbwiddis @gbbafna @kotwanikunal @linuxpi @mch2 @msfroh @nknize @owaiskazi19 @reta @Rishikesh1159 @sachinpkale @saratvemulapalli @shwetathareja @sohami @VachaShah /modules/lang-painless/ @anasalkouz @andrross @ashking94 @Bukhtawar @CEHENKLE @dblock @dbwiddis @gbbafna @jed326 @kotwanikunal @mch2 @msfroh @nknize @owaiskazi19 @reta @Rishikesh1159 @sachinpkale @saratvemulapalli @shwetathareja @sohami @VachaShah /modules/parent-join/ @anasalkouz @andrross @ashking94 @Bukhtawar @CEHENKLE @dblock @dbwiddis @gbbafna @jed326 @kotwanikunal @mch2 @msfroh @nknize @owaiskazi19 @reta @Rishikesh1159 @sachinpkale @saratvemulapalli @shwetathareja @sohami @VachaShah diff --git a/MAINTAINERS.md b/MAINTAINERS.md index f77c69ddeff2a..7be19fbc877d7 100644 --- a/MAINTAINERS.md +++ b/MAINTAINERS.md @@ -5,7 +5,7 @@ This document contains a list of maintainers in this repo. See [opensearch-proje ## Current Maintainers | Maintainer | GitHub ID | Affiliation | -| ------------------------ | ------------------------------------------------------- | ----------- | +|--------------------------|---------------------------------------------------------|-------------| | Anas Alkouz | [anasalkouz](https://github.com/anasalkouz) | Amazon | | Andrew Ross | [andrross](https://github.com/andrross) | Amazon | | Andriy Redko | [reta](https://github.com/reta) | Aiven | @@ -18,6 +18,7 @@ This document contains a list of maintainers in this repo. See [opensearch-proje | Gaurav Bafna | [gbbafna](https://github.com/gbbafna) | Amazon | | Jay Deng | [jed326](https://github.com/jed326) | Amazon | | Kunal Kotwani | [kotwanikunal](https://github.com/kotwanikunal) | Amazon | +| Varun Bansal | [linuxpi](https://github.com/linuxpi) | Amazon | | Marc Handalian | [mch2](https://github.com/mch2) | Amazon | | Michael Froh | [msfroh](https://github.com/msfroh) | Amazon | | Nick Knize | [nknize](https://github.com/nknize) | Amazon | From 4fe03586a95f956b596bfca6c1a5fcd11751072a Mon Sep 17 00:00:00 2001 From: gaobinlong Date: Thu, 8 Aug 2024 19:31:54 +0800 Subject: [PATCH 06/40] Fix yml test {p0=indices.delete_index_template/10_basic} is flaky (#15162) * Fix yml test {p0=indices.delete_index_template/10_basic} is flaky Signed-off-by: Gao Binlong * Fix warning message Signed-off-by: Gao Binlong --------- Signed-off-by: Gao Binlong --- .../test/indices.delete_index_template/10_basic.yml | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/rest-api-spec/src/main/resources/rest-api-spec/test/indices.delete_index_template/10_basic.yml b/rest-api-spec/src/main/resources/rest-api-spec/test/indices.delete_index_template/10_basic.yml index c8c08a2d088ac..6239eb7b8cd22 100644 --- a/rest-api-spec/src/main/resources/rest-api-spec/test/indices.delete_index_template/10_basic.yml +++ b/rest-api-spec/src/main/resources/rest-api-spec/test/indices.delete_index_template/10_basic.yml @@ -1,5 +1,9 @@ setup: + - skip: + features: allowed_warnings - do: + allowed_warnings: + - "index template [test_template_1] has index patterns [test-*] matching patterns from existing older templates [global] with patterns (global => [*]); this template [test_template_1] will take precedence during new index creation" indices.put_index_template: name: test_template_1 body: @@ -11,6 +15,8 @@ setup: "priority": 50 - do: + allowed_warnings: + - "index template [test_template_2] has index patterns [test-*] matching patterns from existing older templates [global] with patterns (global => [*]); this template [test_template_2] will take precedence during new index creation" indices.put_index_template: name: test_template_2 body: From f03dde9c6c280d6af34c7baa9fda6e04d07a0d74 Mon Sep 17 00:00:00 2001 From: "opensearch-trigger-bot[bot]" <98922864+opensearch-trigger-bot[bot]@users.noreply.github.com> Date: Thu, 8 Aug 2024 13:45:54 -0400 Subject: [PATCH 07/40] [AUTO] [main] Add bwc version 2.16.1. (#15160) * Add bwc version 2.16.1 Signed-off-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> * Update Version.java Signed-off-by: Andriy Redko --------- Signed-off-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> Signed-off-by: Andriy Redko Co-authored-by: opensearch-ci-bot <83309141+opensearch-ci-bot@users.noreply.github.com> Co-authored-by: Andriy Redko --- .ci/bwcVersions | 1 + libs/core/src/main/java/org/opensearch/Version.java | 1 + 2 files changed, 2 insertions(+) diff --git a/.ci/bwcVersions b/.ci/bwcVersions index 771bfe694b698..980e08c696d54 100644 --- a/.ci/bwcVersions +++ b/.ci/bwcVersions @@ -36,4 +36,5 @@ BWC_VERSION: - "2.15.0" - "2.15.1" - "2.16.0" + - "2.16.1" - "2.17.0" diff --git a/libs/core/src/main/java/org/opensearch/Version.java b/libs/core/src/main/java/org/opensearch/Version.java index c2d8ce9be29dd..28cb989185ff7 100644 --- a/libs/core/src/main/java/org/opensearch/Version.java +++ b/libs/core/src/main/java/org/opensearch/Version.java @@ -107,6 +107,7 @@ public class Version implements Comparable, ToXContentFragment { public static final Version V_2_15_0 = new Version(2150099, org.apache.lucene.util.Version.LUCENE_9_10_0); public static final Version V_2_15_1 = new Version(2150199, org.apache.lucene.util.Version.LUCENE_9_10_0); public static final Version V_2_16_0 = new Version(2160099, org.apache.lucene.util.Version.LUCENE_9_11_1); + public static final Version V_2_16_1 = new Version(2160199, org.apache.lucene.util.Version.LUCENE_9_11_1); public static final Version V_2_17_0 = new Version(2170099, org.apache.lucene.util.Version.LUCENE_9_11_1); public static final Version V_3_0_0 = new Version(3000099, org.apache.lucene.util.Version.LUCENE_9_12_0); public static final Version CURRENT = V_3_0_0; From 555a56dba55da5be98abbff5ed2fad4d7320f073 Mon Sep 17 00:00:00 2001 From: Rishab Nahata Date: Fri, 9 Aug 2024 00:44:57 +0530 Subject: [PATCH 08/40] Optimise unassigned shards iteration after allocator timeout (#14977) * Optimise unassigned shards iteration after allocator timeout Signed-off-by: Rishab Nahata --- .../common/util/BatchRunnableExecutor.java | 7 ++++ .../gateway/BaseGatewayShardAllocator.java | 11 ++---- .../gateway/ShardsBatchGatewayAllocator.java | 30 ++++++++++----- .../util/BatchRunnableExecutorTests.java | 37 +++++++++++++++++-- .../gateway/GatewayAllocatorTests.java | 21 +++++++++++ .../PrimaryShardBatchAllocatorTests.java | 33 +++++++++++++---- .../ReplicaShardBatchAllocatorTests.java | 8 ++-- .../TestShardBatchGatewayAllocator.java | 21 +++++++++++ 8 files changed, 136 insertions(+), 32 deletions(-) diff --git a/server/src/main/java/org/opensearch/common/util/BatchRunnableExecutor.java b/server/src/main/java/org/opensearch/common/util/BatchRunnableExecutor.java index d3d3304cb909a..cfe2bbb85bda4 100644 --- a/server/src/main/java/org/opensearch/common/util/BatchRunnableExecutor.java +++ b/server/src/main/java/org/opensearch/common/util/BatchRunnableExecutor.java @@ -61,6 +61,13 @@ public void run() { "Time taken to execute timed runnables in this cycle:[{}ms]", TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startTime) ); + onComplete(); } + /** + * Callback method that is invoked after all {@link TimeoutAwareRunnable} instances in the batch have been processed. + * By default, this method does nothing, but it can be overridden by subclasses or modified in the implementation if + * there is a need to perform additional actions once the batch execution is completed. + */ + public void onComplete() {} } diff --git a/server/src/main/java/org/opensearch/gateway/BaseGatewayShardAllocator.java b/server/src/main/java/org/opensearch/gateway/BaseGatewayShardAllocator.java index 0d6af943d39e0..41704545c7a6f 100644 --- a/server/src/main/java/org/opensearch/gateway/BaseGatewayShardAllocator.java +++ b/server/src/main/java/org/opensearch/gateway/BaseGatewayShardAllocator.java @@ -47,7 +47,6 @@ import org.opensearch.core.index.shard.ShardId; import java.util.ArrayList; -import java.util.HashSet; import java.util.List; import java.util.Set; @@ -82,17 +81,15 @@ public void allocateUnassigned( executeDecision(shardRouting, allocateUnassignedDecision, allocation, unassignedAllocationHandler); } - protected void allocateUnassignedBatchOnTimeout(List shardRoutings, RoutingAllocation allocation, boolean primary) { - Set shardIdsFromBatch = new HashSet<>(); - for (ShardRouting shardRouting : shardRoutings) { - ShardId shardId = shardRouting.shardId(); - shardIdsFromBatch.add(shardId); + protected void allocateUnassignedBatchOnTimeout(Set shardIds, RoutingAllocation allocation, boolean primary) { + if (shardIds.isEmpty()) { + return; } RoutingNodes.UnassignedShards.UnassignedIterator iterator = allocation.routingNodes().unassigned().iterator(); while (iterator.hasNext()) { ShardRouting unassignedShard = iterator.next(); AllocateUnassignedDecision allocationDecision; - if (unassignedShard.primary() == primary && shardIdsFromBatch.contains(unassignedShard.shardId())) { + if (unassignedShard.primary() == primary && shardIds.contains(unassignedShard.shardId())) { allocationDecision = AllocateUnassignedDecision.throttle(null); executeDecision(unassignedShard, allocationDecision, allocation, iterator); } diff --git a/server/src/main/java/org/opensearch/gateway/ShardsBatchGatewayAllocator.java b/server/src/main/java/org/opensearch/gateway/ShardsBatchGatewayAllocator.java index 6c6b1126a78d6..d18304ea73ed0 100644 --- a/server/src/main/java/org/opensearch/gateway/ShardsBatchGatewayAllocator.java +++ b/server/src/main/java/org/opensearch/gateway/ShardsBatchGatewayAllocator.java @@ -277,17 +277,14 @@ protected BatchRunnableExecutor innerAllocateUnassignedBatch( } List runnables = new ArrayList<>(); if (primary) { + Set timedOutPrimaryShardIds = new HashSet<>(); batchIdToStartedShardBatch.values() .stream() .filter(batch -> batchesToAssign.contains(batch.batchId)) .forEach(shardsBatch -> runnables.add(new TimeoutAwareRunnable() { @Override public void onTimeout() { - primaryBatchShardAllocator.allocateUnassignedBatchOnTimeout( - shardsBatch.getBatchedShardRoutings(), - allocation, - true - ); + timedOutPrimaryShardIds.addAll(shardsBatch.getBatchedShards()); } @Override @@ -295,15 +292,22 @@ public void run() { primaryBatchShardAllocator.allocateUnassignedBatch(shardsBatch.getBatchedShardRoutings(), allocation); } })); - return new BatchRunnableExecutor(runnables, () -> primaryShardsBatchGatewayAllocatorTimeout); + return new BatchRunnableExecutor(runnables, () -> primaryShardsBatchGatewayAllocatorTimeout) { + @Override + public void onComplete() { + logger.trace("Triggering oncomplete after timeout for [{}] primary shards", timedOutPrimaryShardIds.size()); + primaryBatchShardAllocator.allocateUnassignedBatchOnTimeout(timedOutPrimaryShardIds, allocation, true); + } + }; } else { + Set timedOutReplicaShardIds = new HashSet<>(); batchIdToStoreShardBatch.values() .stream() .filter(batch -> batchesToAssign.contains(batch.batchId)) .forEach(batch -> runnables.add(new TimeoutAwareRunnable() { @Override public void onTimeout() { - replicaBatchShardAllocator.allocateUnassignedBatchOnTimeout(batch.getBatchedShardRoutings(), allocation, false); + timedOutReplicaShardIds.addAll(batch.getBatchedShards()); } @Override @@ -311,7 +315,13 @@ public void run() { replicaBatchShardAllocator.allocateUnassignedBatch(batch.getBatchedShardRoutings(), allocation); } })); - return new BatchRunnableExecutor(runnables, () -> replicaShardsBatchGatewayAllocatorTimeout); + return new BatchRunnableExecutor(runnables, () -> replicaShardsBatchGatewayAllocatorTimeout) { + @Override + public void onComplete() { + logger.trace("Triggering oncomplete after timeout for [{}] replica shards", timedOutReplicaShardIds.size()); + replicaBatchShardAllocator.allocateUnassignedBatchOnTimeout(timedOutReplicaShardIds, allocation, false); + } + }; } } @@ -846,11 +856,11 @@ public int getNumberOfStoreShardBatches() { return batchIdToStoreShardBatch.size(); } - private void setPrimaryBatchAllocatorTimeout(TimeValue primaryShardsBatchGatewayAllocatorTimeout) { + protected void setPrimaryBatchAllocatorTimeout(TimeValue primaryShardsBatchGatewayAllocatorTimeout) { this.primaryShardsBatchGatewayAllocatorTimeout = primaryShardsBatchGatewayAllocatorTimeout; } - private void setReplicaBatchAllocatorTimeout(TimeValue replicaShardsBatchGatewayAllocatorTimeout) { + protected void setReplicaBatchAllocatorTimeout(TimeValue replicaShardsBatchGatewayAllocatorTimeout) { this.replicaShardsBatchGatewayAllocatorTimeout = replicaShardsBatchGatewayAllocatorTimeout; } } diff --git a/server/src/test/java/org/opensearch/common/util/BatchRunnableExecutorTests.java b/server/src/test/java/org/opensearch/common/util/BatchRunnableExecutorTests.java index 269f89faec54d..2f63ae43d0ded 100644 --- a/server/src/test/java/org/opensearch/common/util/BatchRunnableExecutorTests.java +++ b/server/src/test/java/org/opensearch/common/util/BatchRunnableExecutorTests.java @@ -15,6 +15,7 @@ import java.util.Arrays; import java.util.Collections; import java.util.List; +import java.util.concurrent.CountDownLatch; import java.util.function.Supplier; import static org.mockito.Mockito.atMost; @@ -42,7 +43,13 @@ public void setupRunnables() { public void testRunWithoutTimeout() { setupRunnables(); timeoutSupplier = () -> TimeValue.timeValueSeconds(1); - BatchRunnableExecutor executor = new BatchRunnableExecutor(runnableList, timeoutSupplier); + CountDownLatch countDownLatch = new CountDownLatch(1); + BatchRunnableExecutor executor = new BatchRunnableExecutor(runnableList, timeoutSupplier) { + @Override + public void onComplete() { + countDownLatch.countDown(); + } + }; executor.run(); verify(runnable1, times(1)).run(); verify(runnable2, times(1)).run(); @@ -50,12 +57,19 @@ public void testRunWithoutTimeout() { verify(runnable1, never()).onTimeout(); verify(runnable2, never()).onTimeout(); verify(runnable3, never()).onTimeout(); + assertEquals(0, countDownLatch.getCount()); } public void testRunWithTimeout() { setupRunnables(); timeoutSupplier = () -> TimeValue.timeValueNanos(1); - BatchRunnableExecutor executor = new BatchRunnableExecutor(runnableList, timeoutSupplier); + CountDownLatch countDownLatch = new CountDownLatch(1); + BatchRunnableExecutor executor = new BatchRunnableExecutor(runnableList, timeoutSupplier) { + @Override + public void onComplete() { + countDownLatch.countDown(); + } + }; executor.run(); verify(runnable1, times(1)).onTimeout(); verify(runnable2, times(1)).onTimeout(); @@ -63,12 +77,19 @@ public void testRunWithTimeout() { verify(runnable1, never()).run(); verify(runnable2, never()).run(); verify(runnable3, never()).run(); + assertEquals(0, countDownLatch.getCount()); } public void testRunWithPartialTimeout() { setupRunnables(); timeoutSupplier = () -> TimeValue.timeValueMillis(50); - BatchRunnableExecutor executor = new BatchRunnableExecutor(runnableList, timeoutSupplier); + CountDownLatch countDownLatch = new CountDownLatch(1); + BatchRunnableExecutor executor = new BatchRunnableExecutor(runnableList, timeoutSupplier) { + @Override + public void onComplete() { + countDownLatch.countDown(); + } + }; doAnswer(invocation -> { Thread.sleep(100); return null; @@ -81,11 +102,18 @@ public void testRunWithPartialTimeout() { verify(runnable3, atMost(1)).onTimeout(); verify(runnable2, atMost(1)).onTimeout(); verify(runnable3, atMost(1)).onTimeout(); + assertEquals(0, countDownLatch.getCount()); } public void testRunWithEmptyRunnableList() { setupRunnables(); - BatchRunnableExecutor executor = new BatchRunnableExecutor(Collections.emptyList(), timeoutSupplier); + CountDownLatch countDownLatch = new CountDownLatch(1); + BatchRunnableExecutor executor = new BatchRunnableExecutor(Collections.emptyList(), timeoutSupplier) { + @Override + public void onComplete() { + countDownLatch.countDown(); + } + }; executor.run(); verify(runnable1, never()).onTimeout(); verify(runnable2, never()).onTimeout(); @@ -93,5 +121,6 @@ public void testRunWithEmptyRunnableList() { verify(runnable1, never()).run(); verify(runnable2, never()).run(); verify(runnable3, never()).run(); + assertEquals(1, countDownLatch.getCount()); } } diff --git a/server/src/test/java/org/opensearch/gateway/GatewayAllocatorTests.java b/server/src/test/java/org/opensearch/gateway/GatewayAllocatorTests.java index 1596a0b566b28..c7eae77d6deba 100644 --- a/server/src/test/java/org/opensearch/gateway/GatewayAllocatorTests.java +++ b/server/src/test/java/org/opensearch/gateway/GatewayAllocatorTests.java @@ -32,6 +32,7 @@ import org.opensearch.cluster.routing.allocation.decider.AllocationDeciders; import org.opensearch.common.collect.Tuple; import org.opensearch.common.settings.Settings; +import org.opensearch.common.unit.TimeValue; import org.opensearch.common.util.BatchRunnableExecutor; import org.opensearch.common.util.set.Sets; import org.opensearch.core.index.shard.ShardId; @@ -45,6 +46,8 @@ import java.util.List; import java.util.Map; import java.util.Set; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; import java.util.stream.Collectors; import static org.opensearch.gateway.ShardsBatchGatewayAllocator.PRIMARY_BATCH_ALLOCATOR_TIMEOUT_SETTING; @@ -423,6 +426,24 @@ public void testReplicaAllocatorTimeout() { assertEquals(-1, REPLICA_BATCH_ALLOCATOR_TIMEOUT_SETTING.get(build).getMillis()); } + public void testCollectTimedOutShards() throws InterruptedException { + createIndexAndUpdateClusterState(2, 5, 2); + CountDownLatch latch = new CountDownLatch(10); + testShardsBatchGatewayAllocator = new TestShardBatchGatewayAllocator(latch); + testShardsBatchGatewayAllocator.setPrimaryBatchAllocatorTimeout(TimeValue.ZERO); + testShardsBatchGatewayAllocator.setReplicaBatchAllocatorTimeout(TimeValue.ZERO); + BatchRunnableExecutor executor = testShardsBatchGatewayAllocator.allocateAllUnassignedShards(testAllocation, true); + executor.run(); + assertTrue(latch.await(1, TimeUnit.MINUTES)); + latch = new CountDownLatch(10); + testShardsBatchGatewayAllocator = new TestShardBatchGatewayAllocator(latch); + testShardsBatchGatewayAllocator.setPrimaryBatchAllocatorTimeout(TimeValue.ZERO); + testShardsBatchGatewayAllocator.setReplicaBatchAllocatorTimeout(TimeValue.ZERO); + executor = testShardsBatchGatewayAllocator.allocateAllUnassignedShards(testAllocation, false); + executor.run(); + assertTrue(latch.await(1, TimeUnit.MINUTES)); + } + private void createIndexAndUpdateClusterState(int count, int numberOfShards, int numberOfReplicas) { if (count == 0) return; Metadata.Builder metadata = Metadata.builder(); diff --git a/server/src/test/java/org/opensearch/gateway/PrimaryShardBatchAllocatorTests.java b/server/src/test/java/org/opensearch/gateway/PrimaryShardBatchAllocatorTests.java index 270cf465d0f80..48183fed66671 100644 --- a/server/src/test/java/org/opensearch/gateway/PrimaryShardBatchAllocatorTests.java +++ b/server/src/test/java/org/opensearch/gateway/PrimaryShardBatchAllocatorTests.java @@ -41,7 +41,6 @@ import org.junit.Before; import java.util.ArrayList; -import java.util.Arrays; import java.util.Collections; import java.util.HashMap; import java.util.HashSet; @@ -264,8 +263,9 @@ public void testAllocateUnassignedBatchOnTimeoutWithMatchingPrimaryShards() { final RoutingAllocation routingAllocation = routingAllocationWithOnePrimary(allocationDeciders, CLUSTER_RECOVERED, "allocId-0"); ShardRouting shardRouting = routingAllocation.routingTable().getIndicesRouting().get("test").shard(shardId.id()).primaryShard(); - List shardRoutings = Arrays.asList(shardRouting); - batchAllocator.allocateUnassignedBatchOnTimeout(shardRoutings, routingAllocation, true); + Set shardIds = new HashSet<>(); + shardIds.add(shardRouting.shardId()); + batchAllocator.allocateUnassignedBatchOnTimeout(shardIds, routingAllocation, true); List ignoredShards = routingAllocation.routingNodes().unassigned().ignored(); assertEquals(1, ignoredShards.size()); @@ -277,8 +277,7 @@ public void testAllocateUnassignedBatchOnTimeoutWithNoMatchingPrimaryShards() { AllocationDeciders allocationDeciders = randomAllocationDeciders(Settings.builder().build(), clusterSettings, random()); setUpShards(1); final RoutingAllocation routingAllocation = routingAllocationWithOnePrimary(allocationDeciders, CLUSTER_RECOVERED, "allocId-0"); - List shardRoutings = new ArrayList<>(); - batchAllocator.allocateUnassignedBatchOnTimeout(shardRoutings, routingAllocation, true); + batchAllocator.allocateUnassignedBatchOnTimeout(new HashSet<>(), routingAllocation, true); List ignoredShards = routingAllocation.routingNodes().unassigned().ignored(); assertEquals(0, ignoredShards.size()); @@ -296,13 +295,33 @@ public void testAllocateUnassignedBatchOnTimeoutWithNonPrimaryShards() { .shard(shardId.id()) .replicaShards() .get(0); - List shardRoutings = Arrays.asList(shardRouting); - batchAllocator.allocateUnassignedBatchOnTimeout(shardRoutings, routingAllocation, false); + Set shardIds = new HashSet<>(); + shardIds.add(shardRouting.shardId()); + batchAllocator.allocateUnassignedBatchOnTimeout(shardIds, routingAllocation, false); List ignoredShards = routingAllocation.routingNodes().unassigned().ignored(); assertEquals(1, ignoredShards.size()); } + public void testAllocateUnassignedBatchOnTimeoutWithNoShards() { + ClusterSettings clusterSettings = new ClusterSettings(Settings.EMPTY, ClusterSettings.BUILT_IN_CLUSTER_SETTINGS); + AllocationDeciders allocationDeciders = randomAllocationDeciders(Settings.builder().build(), clusterSettings, random()); + setUpShards(1); + final RoutingAllocation routingAllocation = routingAllocationWithOnePrimary(allocationDeciders, CLUSTER_RECOVERED, "allocId-0"); + + ShardRouting shardRouting = routingAllocation.routingTable() + .getIndicesRouting() + .get("test") + .shard(shardId.id()) + .replicaShards() + .get(0); + Set shardIds = new HashSet<>(); + batchAllocator.allocateUnassignedBatchOnTimeout(shardIds, routingAllocation, false); + + List ignoredShards = routingAllocation.routingNodes().unassigned().ignored(); + assertEquals(0, ignoredShards.size()); + } + private RoutingAllocation routingAllocationWithOnePrimary( AllocationDeciders deciders, UnassignedInfo.Reason reason, diff --git a/server/src/test/java/org/opensearch/gateway/ReplicaShardBatchAllocatorTests.java b/server/src/test/java/org/opensearch/gateway/ReplicaShardBatchAllocatorTests.java index 435fd78be2bcd..78ed3f2c7d38c 100644 --- a/server/src/test/java/org/opensearch/gateway/ReplicaShardBatchAllocatorTests.java +++ b/server/src/test/java/org/opensearch/gateway/ReplicaShardBatchAllocatorTests.java @@ -720,9 +720,9 @@ public void testAllocateUnassignedBatchThrottlingAllocationDeciderIsHonoured() t public void testAllocateUnassignedBatchOnTimeoutWithUnassignedReplicaShard() { RoutingAllocation allocation = onePrimaryOnNode1And1Replica(yesAllocationDeciders()); final RoutingNodes.UnassignedShards.UnassignedIterator iterator = allocation.routingNodes().unassigned().iterator(); - List shards = new ArrayList<>(); + Set shards = new HashSet<>(); while (iterator.hasNext()) { - shards.add(iterator.next()); + shards.add(iterator.next().shardId()); } testBatchAllocator.allocateUnassignedBatchOnTimeout(shards, allocation, false); assertThat(allocation.routingNodes().unassigned().ignored().size(), equalTo(1)); @@ -736,9 +736,9 @@ public void testAllocateUnassignedBatchOnTimeoutWithUnassignedReplicaShard() { public void testAllocateUnassignedBatchOnTimeoutWithAlreadyRecoveringReplicaShard() { RoutingAllocation allocation = onePrimaryOnNode1And1ReplicaRecovering(yesAllocationDeciders()); final RoutingNodes.UnassignedShards.UnassignedIterator iterator = allocation.routingNodes().unassigned().iterator(); - List shards = new ArrayList<>(); + Set shards = new HashSet<>(); while (iterator.hasNext()) { - shards.add(iterator.next()); + shards.add(iterator.next().shardId()); } testBatchAllocator.allocateUnassignedBatchOnTimeout(shards, allocation, false); assertThat(allocation.routingNodes().unassigned().ignored().size(), equalTo(0)); diff --git a/test/framework/src/main/java/org/opensearch/test/gateway/TestShardBatchGatewayAllocator.java b/test/framework/src/main/java/org/opensearch/test/gateway/TestShardBatchGatewayAllocator.java index 0eb4bb6935bac..156b1d7c620e6 100644 --- a/test/framework/src/main/java/org/opensearch/test/gateway/TestShardBatchGatewayAllocator.java +++ b/test/framework/src/main/java/org/opensearch/test/gateway/TestShardBatchGatewayAllocator.java @@ -29,13 +29,20 @@ import java.util.List; import java.util.Map; import java.util.Set; +import java.util.concurrent.CountDownLatch; public class TestShardBatchGatewayAllocator extends ShardsBatchGatewayAllocator { + CountDownLatch latch; + public TestShardBatchGatewayAllocator() { } + public TestShardBatchGatewayAllocator(CountDownLatch latch) { + this.latch = latch; + } + public TestShardBatchGatewayAllocator(long maxBatchSize) { super(maxBatchSize); } @@ -83,6 +90,13 @@ protected AsyncShardFetch.FetchResult(foundShards, shardsToIgnoreNodes); } + + @Override + protected void allocateUnassignedBatchOnTimeout(Set shardIds, RoutingAllocation allocation, boolean primary) { + for (int i = 0; i < shardIds.size(); i++) { + latch.countDown(); + } + } }; ReplicaShardBatchAllocator replicaBatchShardAllocator = new ReplicaShardBatchAllocator() { @@ -100,6 +114,13 @@ protected AsyncShardFetch.FetchResult shardIds, RoutingAllocation allocation, boolean primary) { + for (int i = 0; i < shardIds.size(); i++) { + latch.countDown(); + } + } }; @Override From 978d14e867c41b224303b2e77745495336568f79 Mon Sep 17 00:00:00 2001 From: Jay Deng Date: Thu, 8 Aug 2024 15:10:03 -0700 Subject: [PATCH 09/40] Add slice level operation listener methods (#15153) Signed-off-by: Jay Deng --- CHANGELOG.md | 1 + .../index/shard/SearchOperationListener.java | 60 ++++++++++ .../search/internal/ContextIndexSearcher.java | 31 ++++-- .../shard/SearchOperationListenerTests.java | 105 ++++++++++++++++++ .../search/SearchCancellationTests.java | 4 + .../internal/ContextIndexSearcherTests.java | 4 + .../profile/query/QueryProfilerTests.java | 4 + .../search/query/QueryPhaseTests.java | 7 ++ .../search/query/QueryProfilePhaseTests.java | 7 ++ .../aggregations/AggregatorTestCase.java | 4 + 10 files changed, 215 insertions(+), 12 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index f44949bf38511..1964e456acda0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,6 +16,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), - Add `rangeQuery` and `regexpQuery` for `constant_keyword` field type ([#14711](https://github.com/opensearch-project/OpenSearch/pull/14711)) - Add took time to request nodes stats ([#15054](https://github.com/opensearch-project/OpenSearch/pull/15054)) - [Workload Management] QueryGroup resource tracking framework changes ([#13897](https://github.com/opensearch-project/OpenSearch/pull/13897)) +- Add slice execution listeners to SearchOperationListener interface ([#15153](https://github.com/opensearch-project/OpenSearch/pull/15153)) ### Dependencies - Bump `netty` from 4.1.111.Final to 4.1.112.Final ([#15081](https://github.com/opensearch-project/OpenSearch/pull/15081)) diff --git a/server/src/main/java/org/opensearch/index/shard/SearchOperationListener.java b/server/src/main/java/org/opensearch/index/shard/SearchOperationListener.java index 94079db468f9c..5e63262df0d70 100644 --- a/server/src/main/java/org/opensearch/index/shard/SearchOperationListener.java +++ b/server/src/main/java/org/opensearch/index/shard/SearchOperationListener.java @@ -71,6 +71,33 @@ default void onFailedQueryPhase(SearchContext searchContext) {} */ default void onQueryPhase(SearchContext searchContext, long tookInNanos) {} + /** + * Executed before the slice execution in + * {@link org.opensearch.search.internal.ContextIndexSearcher#search(List, org.apache.lucene.search.Weight, org.apache.lucene.search.Collector)}. + * This will be called once per slice in concurrent search and only once in non-concurrent search. + * @param searchContext the current search context + */ + default void onPreSliceExecution(SearchContext searchContext) {} + + /** + * Executed if the slice execution in + * {@link org.opensearch.search.internal.ContextIndexSearcher#search(List, org.apache.lucene.search.Weight, org.apache.lucene.search.Collector)} failed. + * This will be called once per slice in concurrent search and only once in non-concurrent search. + * @param searchContext the current search context + */ + default void onFailedSliceExecution(SearchContext searchContext) {} + + /** + * Executed after the slice execution in + * {@link org.opensearch.search.internal.ContextIndexSearcher#search(List, org.apache.lucene.search.Weight, org.apache.lucene.search.Collector)} successfully finished. + * This will be called once per slice in concurrent search and only once in non-concurrent search. + * Note: this is not invoked if the slice execution failed.* + * @param searchContext the current search context + * + * @see #onFailedSliceExecution(org.opensearch.search.internal.SearchContext) + */ + default void onSliceExecution(SearchContext searchContext) {} + /** * Executed before the fetch phase is executed * @param searchContext the current search context @@ -195,6 +222,39 @@ public void onQueryPhase(SearchContext searchContext, long tookInNanos) { } } + @Override + public void onPreSliceExecution(SearchContext searchContext) { + for (SearchOperationListener listener : listeners) { + try { + listener.onPreSliceExecution(searchContext); + } catch (Exception e) { + logger.warn(() -> new ParameterizedMessage("onPreSliceExecution listener [{}] failed", listener), e); + } + } + } + + @Override + public void onFailedSliceExecution(SearchContext searchContext) { + for (SearchOperationListener listener : listeners) { + try { + listener.onFailedSliceExecution(searchContext); + } catch (Exception e) { + logger.warn(() -> new ParameterizedMessage("onFailedSliceExecution listener [{}] failed", listener), e); + } + } + } + + @Override + public void onSliceExecution(SearchContext searchContext) { + for (SearchOperationListener listener : listeners) { + try { + listener.onSliceExecution(searchContext); + } catch (Exception e) { + logger.warn(() -> new ParameterizedMessage("onSliceExecution listener [{}] failed", listener), e); + } + } + } + @Override public void onPreFetchPhase(SearchContext searchContext) { for (SearchOperationListener listener : listeners) { diff --git a/server/src/main/java/org/opensearch/search/internal/ContextIndexSearcher.java b/server/src/main/java/org/opensearch/search/internal/ContextIndexSearcher.java index ec3ed2332d0b8..fa00ace378df1 100644 --- a/server/src/main/java/org/opensearch/search/internal/ContextIndexSearcher.java +++ b/server/src/main/java/org/opensearch/search/internal/ContextIndexSearcher.java @@ -270,20 +270,27 @@ public void search( @Override protected void search(List leaves, Weight weight, Collector collector) throws IOException { - // Time series based workload by default traverses segments in desc order i.e. latest to the oldest order. - // This is actually beneficial for search queries to start search on latest segments first for time series workload. - // That can slow down ASC order queries on timestamp workload. So to avoid that slowdown, we will reverse leaf - // reader order here. - if (searchContext.shouldUseTimeSeriesDescSortOptimization()) { - for (int i = leaves.size() - 1; i >= 0; i--) { - searchLeaf(leaves.get(i), weight, collector); - } - } else { - for (int i = 0; i < leaves.size(); i++) { - searchLeaf(leaves.get(i), weight, collector); + searchContext.indexShard().getSearchOperationListener().onPreSliceExecution(searchContext); + try { + // Time series based workload by default traverses segments in desc order i.e. latest to the oldest order. + // This is actually beneficial for search queries to start search on latest segments first for time series workload. + // That can slow down ASC order queries on timestamp workload. So to avoid that slowdown, we will reverse leaf + // reader order here. + if (searchContext.shouldUseTimeSeriesDescSortOptimization()) { + for (int i = leaves.size() - 1; i >= 0; i--) { + searchLeaf(leaves.get(i), weight, collector); + } + } else { + for (int i = 0; i < leaves.size(); i++) { + searchLeaf(leaves.get(i), weight, collector); + } } + searchContext.bucketCollectorProcessor().processPostCollection(collector); + } catch (Throwable t) { + searchContext.indexShard().getSearchOperationListener().onFailedSliceExecution(searchContext); + throw t; } - searchContext.bucketCollectorProcessor().processPostCollection(collector); + searchContext.indexShard().getSearchOperationListener().onSliceExecution(searchContext); } /** diff --git a/server/src/test/java/org/opensearch/index/shard/SearchOperationListenerTests.java b/server/src/test/java/org/opensearch/index/shard/SearchOperationListenerTests.java index c61c13eecf2c3..b00307920e875 100644 --- a/server/src/test/java/org/opensearch/index/shard/SearchOperationListenerTests.java +++ b/server/src/test/java/org/opensearch/index/shard/SearchOperationListenerTests.java @@ -56,6 +56,9 @@ public void testListenersAreExecuted() { AtomicInteger preQuery = new AtomicInteger(); AtomicInteger failedQuery = new AtomicInteger(); AtomicInteger onQuery = new AtomicInteger(); + AtomicInteger preSlice = new AtomicInteger(); + AtomicInteger failedSlice = new AtomicInteger(); + AtomicInteger onSlice = new AtomicInteger(); AtomicInteger onFetch = new AtomicInteger(); AtomicInteger preFetch = new AtomicInteger(); AtomicInteger failedFetch = new AtomicInteger(); @@ -86,6 +89,24 @@ public void onQueryPhase(SearchContext searchContext, long tookInNanos) { onQuery.incrementAndGet(); } + @Override + public void onPreSliceExecution(SearchContext searchContext) { + assertNotNull(searchContext); + preSlice.incrementAndGet(); + } + + @Override + public void onFailedSliceExecution(SearchContext searchContext) { + assertNotNull(searchContext); + failedSlice.incrementAndGet(); + } + + @Override + public void onSliceExecution(SearchContext searchContext) { + assertNotNull(searchContext); + onSlice.incrementAndGet(); + } + @Override public void onPreFetchPhase(SearchContext searchContext) { assertNotNull(searchContext); @@ -167,10 +188,30 @@ public void onSearchIdleReactivation() { compositeListener.onQueryPhase(ctx, timeInNanos.get()); assertEquals(0, preFetch.get()); assertEquals(0, preQuery.get()); + assertEquals(0, preSlice.get()); + assertEquals(0, failedFetch.get()); + assertEquals(0, failedQuery.get()); + assertEquals(0, failedSlice.get()); + assertEquals(2, onQuery.get()); + assertEquals(0, onFetch.get()); + assertEquals(0, onSlice.get()); + assertEquals(0, newContext.get()); + assertEquals(0, newScrollContext.get()); + assertEquals(0, freeContext.get()); + assertEquals(0, freeScrollContext.get()); + assertEquals(0, searchIdleReactivateCount.get()); + assertEquals(0, validateSearchContext.get()); + + compositeListener.onSliceExecution(ctx); + assertEquals(0, preFetch.get()); + assertEquals(0, preQuery.get()); + assertEquals(0, preSlice.get()); assertEquals(0, failedFetch.get()); assertEquals(0, failedQuery.get()); + assertEquals(0, failedSlice.get()); assertEquals(2, onQuery.get()); assertEquals(0, onFetch.get()); + assertEquals(2, onSlice.get()); assertEquals(0, newContext.get()); assertEquals(0, newScrollContext.get()); assertEquals(0, freeContext.get()); @@ -181,10 +222,13 @@ public void onSearchIdleReactivation() { compositeListener.onFetchPhase(ctx, timeInNanos.get()); assertEquals(0, preFetch.get()); assertEquals(0, preQuery.get()); + assertEquals(0, preSlice.get()); assertEquals(0, failedFetch.get()); assertEquals(0, failedQuery.get()); + assertEquals(0, failedSlice.get()); assertEquals(2, onQuery.get()); assertEquals(2, onFetch.get()); + assertEquals(2, onSlice.get()); assertEquals(0, newContext.get()); assertEquals(0, newScrollContext.get()); assertEquals(0, freeContext.get()); @@ -195,10 +239,30 @@ public void onSearchIdleReactivation() { compositeListener.onPreQueryPhase(ctx); assertEquals(0, preFetch.get()); assertEquals(2, preQuery.get()); + assertEquals(0, preSlice.get()); assertEquals(0, failedFetch.get()); assertEquals(0, failedQuery.get()); + assertEquals(0, failedSlice.get()); assertEquals(2, onQuery.get()); assertEquals(2, onFetch.get()); + assertEquals(2, onSlice.get()); + assertEquals(0, newContext.get()); + assertEquals(0, newScrollContext.get()); + assertEquals(0, freeContext.get()); + assertEquals(0, freeScrollContext.get()); + assertEquals(0, searchIdleReactivateCount.get()); + assertEquals(0, validateSearchContext.get()); + + compositeListener.onPreSliceExecution(ctx); + assertEquals(0, preFetch.get()); + assertEquals(2, preQuery.get()); + assertEquals(2, preSlice.get()); + assertEquals(0, failedFetch.get()); + assertEquals(0, failedQuery.get()); + assertEquals(0, failedSlice.get()); + assertEquals(2, onQuery.get()); + assertEquals(2, onFetch.get()); + assertEquals(2, onSlice.get()); assertEquals(0, newContext.get()); assertEquals(0, newScrollContext.get()); assertEquals(0, freeContext.get()); @@ -209,10 +273,13 @@ public void onSearchIdleReactivation() { compositeListener.onPreFetchPhase(ctx); assertEquals(2, preFetch.get()); assertEquals(2, preQuery.get()); + assertEquals(2, preSlice.get()); assertEquals(0, failedFetch.get()); assertEquals(0, failedQuery.get()); + assertEquals(0, failedSlice.get()); assertEquals(2, onQuery.get()); assertEquals(2, onFetch.get()); + assertEquals(2, onSlice.get()); assertEquals(0, newContext.get()); assertEquals(0, newScrollContext.get()); assertEquals(0, freeContext.get()); @@ -223,10 +290,13 @@ public void onSearchIdleReactivation() { compositeListener.onFailedFetchPhase(ctx); assertEquals(2, preFetch.get()); assertEquals(2, preQuery.get()); + assertEquals(2, preSlice.get()); assertEquals(2, failedFetch.get()); assertEquals(0, failedQuery.get()); + assertEquals(0, failedSlice.get()); assertEquals(2, onQuery.get()); assertEquals(2, onFetch.get()); + assertEquals(2, onSlice.get()); assertEquals(0, newContext.get()); assertEquals(0, newScrollContext.get()); assertEquals(0, freeContext.get()); @@ -237,10 +307,30 @@ public void onSearchIdleReactivation() { compositeListener.onFailedQueryPhase(ctx); assertEquals(2, preFetch.get()); assertEquals(2, preQuery.get()); + assertEquals(2, preSlice.get()); + assertEquals(2, failedFetch.get()); + assertEquals(2, failedQuery.get()); + assertEquals(0, failedSlice.get()); + assertEquals(2, onQuery.get()); + assertEquals(2, onFetch.get()); + assertEquals(2, onSlice.get()); + assertEquals(0, newContext.get()); + assertEquals(0, newScrollContext.get()); + assertEquals(0, freeContext.get()); + assertEquals(0, freeScrollContext.get()); + assertEquals(0, searchIdleReactivateCount.get()); + assertEquals(0, validateSearchContext.get()); + + compositeListener.onFailedSliceExecution(ctx); + assertEquals(2, preFetch.get()); + assertEquals(2, preQuery.get()); + assertEquals(2, preSlice.get()); assertEquals(2, failedFetch.get()); assertEquals(2, failedQuery.get()); + assertEquals(2, failedSlice.get()); assertEquals(2, onQuery.get()); assertEquals(2, onFetch.get()); + assertEquals(2, onSlice.get()); assertEquals(0, newContext.get()); assertEquals(0, newScrollContext.get()); assertEquals(0, freeContext.get()); @@ -251,10 +341,13 @@ public void onSearchIdleReactivation() { compositeListener.onNewReaderContext(mock(ReaderContext.class)); assertEquals(2, preFetch.get()); assertEquals(2, preQuery.get()); + assertEquals(2, preSlice.get()); assertEquals(2, failedFetch.get()); assertEquals(2, failedQuery.get()); + assertEquals(2, failedSlice.get()); assertEquals(2, onQuery.get()); assertEquals(2, onFetch.get()); + assertEquals(2, onSlice.get()); assertEquals(2, newContext.get()); assertEquals(0, newScrollContext.get()); assertEquals(0, freeContext.get()); @@ -265,10 +358,13 @@ public void onSearchIdleReactivation() { compositeListener.onNewScrollContext(mock(ReaderContext.class)); assertEquals(2, preFetch.get()); assertEquals(2, preQuery.get()); + assertEquals(2, preSlice.get()); assertEquals(2, failedFetch.get()); assertEquals(2, failedQuery.get()); + assertEquals(2, failedSlice.get()); assertEquals(2, onQuery.get()); assertEquals(2, onFetch.get()); + assertEquals(2, onSlice.get()); assertEquals(2, newContext.get()); assertEquals(2, newScrollContext.get()); assertEquals(0, freeContext.get()); @@ -279,10 +375,13 @@ public void onSearchIdleReactivation() { compositeListener.onFreeReaderContext(mock(ReaderContext.class)); assertEquals(2, preFetch.get()); assertEquals(2, preQuery.get()); + assertEquals(2, preSlice.get()); assertEquals(2, failedFetch.get()); assertEquals(2, failedQuery.get()); + assertEquals(2, failedSlice.get()); assertEquals(2, onQuery.get()); assertEquals(2, onFetch.get()); + assertEquals(2, onSlice.get()); assertEquals(2, newContext.get()); assertEquals(2, newScrollContext.get()); assertEquals(2, freeContext.get()); @@ -293,10 +392,13 @@ public void onSearchIdleReactivation() { compositeListener.onFreeScrollContext(mock(ReaderContext.class)); assertEquals(2, preFetch.get()); assertEquals(2, preQuery.get()); + assertEquals(2, preSlice.get()); assertEquals(2, failedFetch.get()); assertEquals(2, failedQuery.get()); + assertEquals(2, failedSlice.get()); assertEquals(2, onQuery.get()); assertEquals(2, onFetch.get()); + assertEquals(2, onSlice.get()); assertEquals(2, newContext.get()); assertEquals(2, newScrollContext.get()); assertEquals(2, freeContext.get()); @@ -307,10 +409,13 @@ public void onSearchIdleReactivation() { compositeListener.onSearchIdleReactivation(); assertEquals(2, preFetch.get()); assertEquals(2, preQuery.get()); + assertEquals(2, preSlice.get()); assertEquals(2, failedFetch.get()); assertEquals(2, failedQuery.get()); + assertEquals(2, failedSlice.get()); assertEquals(2, onQuery.get()); assertEquals(2, onFetch.get()); + assertEquals(2, onSlice.get()); assertEquals(2, newContext.get()); assertEquals(2, newScrollContext.get()); assertEquals(2, freeContext.get()); diff --git a/server/src/test/java/org/opensearch/search/SearchCancellationTests.java b/server/src/test/java/org/opensearch/search/SearchCancellationTests.java index fce58eecbafb1..266052444da46 100644 --- a/server/src/test/java/org/opensearch/search/SearchCancellationTests.java +++ b/server/src/test/java/org/opensearch/search/SearchCancellationTests.java @@ -51,6 +51,7 @@ import org.opensearch.common.util.io.IOUtils; import org.opensearch.core.tasks.TaskCancelledException; import org.opensearch.index.shard.IndexShard; +import org.opensearch.index.shard.SearchOperationListener; import org.opensearch.search.internal.ContextIndexSearcher; import org.opensearch.search.internal.SearchContext; import org.opensearch.test.OpenSearchTestCase; @@ -138,6 +139,9 @@ public void testCancellableCollector() throws IOException { Runnable cancellation = () -> { throw new TaskCancelledException("cancelled"); }; IndexShard indexShard = mock(IndexShard.class); when(searchContext.indexShard()).thenReturn(indexShard); + SearchOperationListener searchOperationListener = new SearchOperationListener() { + }; + when(indexShard.getSearchOperationListener()).thenReturn(searchOperationListener); ContextIndexSearcher searcher = new ContextIndexSearcher( reader, IndexSearcher.getDefaultSimilarity(), diff --git a/server/src/test/java/org/opensearch/search/internal/ContextIndexSearcherTests.java b/server/src/test/java/org/opensearch/search/internal/ContextIndexSearcherTests.java index a707c8b34e0a4..606c2512a3d58 100644 --- a/server/src/test/java/org/opensearch/search/internal/ContextIndexSearcherTests.java +++ b/server/src/test/java/org/opensearch/search/internal/ContextIndexSearcherTests.java @@ -81,6 +81,7 @@ import org.opensearch.index.IndexSettings; import org.opensearch.index.cache.bitset.BitsetFilterCache; import org.opensearch.index.shard.IndexShard; +import org.opensearch.index.shard.SearchOperationListener; import org.opensearch.search.SearchService; import org.opensearch.search.aggregations.LeafBucketCollector; import org.opensearch.test.IndexSettingsModule; @@ -262,6 +263,9 @@ public void onRemoval(ShardId shardId, Accountable accountable) { SearchContext searchContext = mock(SearchContext.class); IndexShard indexShard = mock(IndexShard.class); when(searchContext.indexShard()).thenReturn(indexShard); + SearchOperationListener searchOperationListener = new SearchOperationListener() { + }; + when(indexShard.getSearchOperationListener()).thenReturn(searchOperationListener); when(searchContext.bucketCollectorProcessor()).thenReturn(SearchContext.NO_OP_BUCKET_COLLECTOR_PROCESSOR); ContextIndexSearcher searcher = new ContextIndexSearcher( filteredReader, diff --git a/server/src/test/java/org/opensearch/search/profile/query/QueryProfilerTests.java b/server/src/test/java/org/opensearch/search/profile/query/QueryProfilerTests.java index 481a224f2ff0e..3a7c711d324c4 100644 --- a/server/src/test/java/org/opensearch/search/profile/query/QueryProfilerTests.java +++ b/server/src/test/java/org/opensearch/search/profile/query/QueryProfilerTests.java @@ -62,6 +62,7 @@ import org.apache.lucene.tests.util.TestUtil; import org.opensearch.common.util.io.IOUtils; import org.opensearch.index.shard.IndexShard; +import org.opensearch.index.shard.SearchOperationListener; import org.opensearch.search.internal.ContextIndexSearcher; import org.opensearch.search.internal.SearchContext; import org.opensearch.search.profile.ProfileResult; @@ -128,6 +129,9 @@ public void setUp() throws Exception { SearchContext searchContext = mock(SearchContext.class); IndexShard indexShard = mock(IndexShard.class); when(searchContext.indexShard()).thenReturn(indexShard); + SearchOperationListener searchOperationListener = new SearchOperationListener() { + }; + when(indexShard.getSearchOperationListener()).thenReturn(searchOperationListener); when(searchContext.bucketCollectorProcessor()).thenReturn(SearchContext.NO_OP_BUCKET_COLLECTOR_PROCESSOR); searcher = new ContextIndexSearcher( reader, diff --git a/server/src/test/java/org/opensearch/search/query/QueryPhaseTests.java b/server/src/test/java/org/opensearch/search/query/QueryPhaseTests.java index 4bd4d406e4391..84057ab1a1b15 100644 --- a/server/src/test/java/org/opensearch/search/query/QueryPhaseTests.java +++ b/server/src/test/java/org/opensearch/search/query/QueryPhaseTests.java @@ -101,6 +101,7 @@ import org.opensearch.index.search.OpenSearchToParentBlockJoinQuery; import org.opensearch.index.shard.IndexShard; import org.opensearch.index.shard.IndexShardTestCase; +import org.opensearch.index.shard.SearchOperationListener; import org.opensearch.lucene.queries.MinDocQuery; import org.opensearch.search.DocValueFormat; import org.opensearch.search.collapse.CollapseBuilder; @@ -1225,6 +1226,9 @@ private static ContextIndexSearcher newContextSearcher(IndexReader reader, Execu SearchContext searchContext = mock(SearchContext.class); IndexShard indexShard = mock(IndexShard.class); when(searchContext.indexShard()).thenReturn(indexShard); + SearchOperationListener searchOperationListener = new SearchOperationListener() { + }; + when(indexShard.getSearchOperationListener()).thenReturn(searchOperationListener); when(searchContext.bucketCollectorProcessor()).thenReturn(SearchContext.NO_OP_BUCKET_COLLECTOR_PROCESSOR); when(searchContext.shouldUseConcurrentSearch()).thenReturn(executor != null); if (executor != null) { @@ -1248,6 +1252,9 @@ private static ContextIndexSearcher newEarlyTerminationContextSearcher(IndexRead SearchContext searchContext = mock(SearchContext.class); IndexShard indexShard = mock(IndexShard.class); when(searchContext.indexShard()).thenReturn(indexShard); + SearchOperationListener searchOperationListener = new SearchOperationListener() { + }; + when(indexShard.getSearchOperationListener()).thenReturn(searchOperationListener); when(searchContext.bucketCollectorProcessor()).thenReturn(SearchContext.NO_OP_BUCKET_COLLECTOR_PROCESSOR); when(searchContext.shouldUseConcurrentSearch()).thenReturn(executor != null); if (executor != null) { diff --git a/server/src/test/java/org/opensearch/search/query/QueryProfilePhaseTests.java b/server/src/test/java/org/opensearch/search/query/QueryProfilePhaseTests.java index 1d545cea67207..a44b654e462f5 100644 --- a/server/src/test/java/org/opensearch/search/query/QueryProfilePhaseTests.java +++ b/server/src/test/java/org/opensearch/search/query/QueryProfilePhaseTests.java @@ -63,6 +63,7 @@ import org.opensearch.index.query.SourceFieldMatchQuery; import org.opensearch.index.shard.IndexShard; import org.opensearch.index.shard.IndexShardTestCase; +import org.opensearch.index.shard.SearchOperationListener; import org.opensearch.lucene.queries.MinDocQuery; import org.opensearch.search.DocValueFormat; import org.opensearch.search.collapse.CollapseBuilder; @@ -1676,6 +1677,9 @@ private static ContextIndexSearcher newContextSearcher(IndexReader reader, Execu SearchContext searchContext = mock(SearchContext.class); IndexShard indexShard = mock(IndexShard.class); when(searchContext.indexShard()).thenReturn(indexShard); + SearchOperationListener searchOperationListener = new SearchOperationListener() { + }; + when(indexShard.getSearchOperationListener()).thenReturn(searchOperationListener); when(searchContext.bucketCollectorProcessor()).thenReturn(SearchContext.NO_OP_BUCKET_COLLECTOR_PROCESSOR); return new ContextIndexSearcher( reader, @@ -1693,6 +1697,9 @@ private static ContextIndexSearcher newEarlyTerminationContextSearcher(IndexRead SearchContext searchContext = mock(SearchContext.class); IndexShard indexShard = mock(IndexShard.class); when(searchContext.indexShard()).thenReturn(indexShard); + SearchOperationListener searchOperationListener = new SearchOperationListener() { + }; + when(indexShard.getSearchOperationListener()).thenReturn(searchOperationListener); when(searchContext.bucketCollectorProcessor()).thenReturn(SearchContext.NO_OP_BUCKET_COLLECTOR_PROCESSOR); return new ContextIndexSearcher( reader, diff --git a/test/framework/src/main/java/org/opensearch/search/aggregations/AggregatorTestCase.java b/test/framework/src/main/java/org/opensearch/search/aggregations/AggregatorTestCase.java index 544fb100a17bf..4abd7fbea9cff 100644 --- a/test/framework/src/main/java/org/opensearch/search/aggregations/AggregatorTestCase.java +++ b/test/framework/src/main/java/org/opensearch/search/aggregations/AggregatorTestCase.java @@ -119,6 +119,7 @@ import org.opensearch.index.mapper.TextFieldMapper; import org.opensearch.index.query.QueryShardContext; import org.opensearch.index.shard.IndexShard; +import org.opensearch.index.shard.SearchOperationListener; import org.opensearch.indices.IndicesModule; import org.opensearch.indices.fielddata.cache.IndicesFieldDataCache; import org.opensearch.indices.mapper.MapperRegistry; @@ -380,6 +381,9 @@ public boolean shouldCache(Query query) { IndexShard indexShard = mock(IndexShard.class); when(indexShard.shardId()).thenReturn(new ShardId("test", "test", 0)); when(searchContext.indexShard()).thenReturn(indexShard); + SearchOperationListener searchOperationListener = new SearchOperationListener() { + }; + when(indexShard.getSearchOperationListener()).thenReturn(searchOperationListener); when(searchContext.aggregations()).thenReturn(new SearchContextAggregations(AggregatorFactories.EMPTY, bucketConsumer)); when(searchContext.query()).thenReturn(query); when(searchContext.bucketCollectorProcessor()).thenReturn(new BucketCollectorProcessor()); From 170ea27a173c8a126bb2f031da9d3504eeb82208 Mon Sep 17 00:00:00 2001 From: bowenlan-amzn Date: Thu, 8 Aug 2024 21:44:00 -0700 Subject: [PATCH 10/40] Refactor the filter rewrite optimization (#14464) * Refactor Split the single Helper classes and move the classes into a new package for any optimization we introduced for search path. Rename the class name to make it more straightforward and general Signed-off-by: bowenlan-amzn * Refactor refactor the canOptimize logic sort out the basic rule about how to provide data from aggregator, and where to put common logic Signed-off-by: bowenlan-amzn * Refactor refactor the data provider and try optimize logic Signed-off-by: bowenlan-amzn * Refactor Signed-off-by: bowenlan-amzn * Refactor extract segment match all logic Signed-off-by: bowenlan-amzn * Refactor Signed-off-by: bowenlan-amzn * Refactor inline class Signed-off-by: bowenlan-amzn * Fix a bug Signed-off-by: bowenlan-amzn * address comment Signed-off-by: bowenlan-amzn * prepareFromSegment now doesn't return Ranges Signed-off-by: bowenlan-amzn * how it looks like when introduce interfaces Signed-off-by: bowenlan-amzn * remove interface, clean up Signed-off-by: bowenlan-amzn * improve doc Signed-off-by: bowenlan-amzn * move multirangetraversal logic to helper Signed-off-by: bowenlan-amzn * improve the refactor package name -> filterrewrite move tree traversal logic to new class add documentation for important abstract methods add sub class for composite aggregation bridge Signed-off-by: bowenlan-amzn * Address Marc's comments Signed-off-by: bowenlan-amzn * Address concurrent segment search concern To save the ranges per segment, now change to a map that save ranges for segments separately. The increment document function "incrementBucketDocCount" should already be thread safe, as it's the same method used by normal aggregation execution path Signed-off-by: bowenlan-amzn * remove circular dependency Signed-off-by: bowenlan-amzn * Address comment - remove map of segment ranges, pass in by calling getRanges when needed - use AtomicInteger for the debug info Signed-off-by: bowenlan-amzn --------- Signed-off-by: bowenlan-amzn --- .../bucket/FastFilterRewriteHelper.java | 849 ------------------ .../bucket/composite/CompositeAggregator.java | 121 ++- .../filterrewrite/AggregatorBridge.java | 84 ++ .../CompositeAggregatorBridge.java | 36 + .../DateHistogramAggregatorBridge.java | 174 ++++ .../FilterRewriteOptimizationContext.java | 189 ++++ .../bucket/filterrewrite/Helper.java | 213 +++++ .../filterrewrite/PointTreeTraversal.java | 223 +++++ .../filterrewrite/RangeAggregatorBridge.java | 96 ++ .../bucket/filterrewrite/Ranges.java | 57 ++ .../bucket/filterrewrite/package-info.java | 19 + .../AutoDateHistogramAggregator.java | 103 +-- .../histogram/DateHistogramAggregator.java | 76 +- .../bucket/range/RangeAggregator.java | 43 +- 14 files changed, 1256 insertions(+), 1027 deletions(-) delete mode 100644 server/src/main/java/org/opensearch/search/aggregations/bucket/FastFilterRewriteHelper.java create mode 100644 server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/AggregatorBridge.java create mode 100644 server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/CompositeAggregatorBridge.java create mode 100644 server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/DateHistogramAggregatorBridge.java create mode 100644 server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/FilterRewriteOptimizationContext.java create mode 100644 server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/Helper.java create mode 100644 server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/PointTreeTraversal.java create mode 100644 server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/RangeAggregatorBridge.java create mode 100644 server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/Ranges.java create mode 100644 server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/package-info.java diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/FastFilterRewriteHelper.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/FastFilterRewriteHelper.java deleted file mode 100644 index 2ab003fb94e33..0000000000000 --- a/server/src/main/java/org/opensearch/search/aggregations/bucket/FastFilterRewriteHelper.java +++ /dev/null @@ -1,849 +0,0 @@ -/* - * SPDX-License-Identifier: Apache-2.0 - * - * The OpenSearch Contributors require contributions made to - * this file be licensed under the Apache-2.0 license or a - * compatible open source license. - */ - -package org.opensearch.search.aggregations.bucket; - -import org.apache.logging.log4j.LogManager; -import org.apache.logging.log4j.Logger; -import org.apache.lucene.document.LongPoint; -import org.apache.lucene.index.DocValues; -import org.apache.lucene.index.LeafReaderContext; -import org.apache.lucene.index.NumericDocValues; -import org.apache.lucene.index.PointValues; -import org.apache.lucene.search.CollectionTerminatedException; -import org.apache.lucene.search.ConstantScoreQuery; -import org.apache.lucene.search.DocIdSetIterator; -import org.apache.lucene.search.FieldExistsQuery; -import org.apache.lucene.search.IndexOrDocValuesQuery; -import org.apache.lucene.search.MatchAllDocsQuery; -import org.apache.lucene.search.PointRangeQuery; -import org.apache.lucene.search.Query; -import org.apache.lucene.search.ScoreMode; -import org.apache.lucene.search.Weight; -import org.apache.lucene.util.ArrayUtil; -import org.apache.lucene.util.NumericUtils; -import org.opensearch.common.CheckedRunnable; -import org.opensearch.common.Rounding; -import org.opensearch.common.lucene.search.function.FunctionScoreQuery; -import org.opensearch.index.mapper.DateFieldMapper; -import org.opensearch.index.mapper.DocCountFieldMapper; -import org.opensearch.index.mapper.MappedFieldType; -import org.opensearch.index.mapper.NumericPointEncoder; -import org.opensearch.index.query.DateRangeIncludingNowQuery; -import org.opensearch.search.aggregations.bucket.composite.CompositeAggregator; -import org.opensearch.search.aggregations.bucket.composite.CompositeValuesSourceConfig; -import org.opensearch.search.aggregations.bucket.composite.RoundingValuesSource; -import org.opensearch.search.aggregations.bucket.histogram.LongBounds; -import org.opensearch.search.aggregations.bucket.range.RangeAggregator.Range; -import org.opensearch.search.aggregations.support.ValuesSource; -import org.opensearch.search.aggregations.support.ValuesSourceConfig; -import org.opensearch.search.internal.SearchContext; - -import java.io.IOException; -import java.util.HashMap; -import java.util.List; -import java.util.Map; -import java.util.OptionalLong; -import java.util.function.BiConsumer; -import java.util.function.Function; - -import static org.opensearch.index.mapper.NumberFieldMapper.NumberType.LONG; -import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; - -/** - * Utility class to help rewrite aggregations into filters. - * Instead of aggregation collects documents one by one, filter may count all documents that match in one pass. - *

- * Currently supported rewrite: - *

    - *
  • date histogram : date range filter. - * Applied: DateHistogramAggregator, AutoDateHistogramAggregator, CompositeAggregator
    • - *
- * - * @opensearch.internal - */ -public final class FastFilterRewriteHelper { - - private FastFilterRewriteHelper() {} - - private static final Logger logger = LogManager.getLogger(FastFilterRewriteHelper.class); - - private static final Map, Function> queryWrappers; - - // Initialize the wrapper map for unwrapping the query - static { - queryWrappers = new HashMap<>(); - queryWrappers.put(ConstantScoreQuery.class, q -> ((ConstantScoreQuery) q).getQuery()); - queryWrappers.put(FunctionScoreQuery.class, q -> ((FunctionScoreQuery) q).getSubQuery()); - queryWrappers.put(DateRangeIncludingNowQuery.class, q -> ((DateRangeIncludingNowQuery) q).getQuery()); - queryWrappers.put(IndexOrDocValuesQuery.class, q -> ((IndexOrDocValuesQuery) q).getIndexQuery()); - } - - /** - * Recursively unwraps query into the concrete form - * for applying the optimization - */ - private static Query unwrapIntoConcreteQuery(Query query) { - while (queryWrappers.containsKey(query.getClass())) { - query = queryWrappers.get(query.getClass()).apply(query); - } - - return query; - } - - /** - * Finds the global min and max bounds of the field for the shard across all segments - * - * @return null if the field is empty or not indexed - */ - private static long[] getShardBounds(final SearchContext context, final String fieldName) throws IOException { - final List leaves = context.searcher().getIndexReader().leaves(); - long min = Long.MAX_VALUE, max = Long.MIN_VALUE; - for (LeafReaderContext leaf : leaves) { - final PointValues values = leaf.reader().getPointValues(fieldName); - if (values != null) { - min = Math.min(min, NumericUtils.sortableBytesToLong(values.getMinPackedValue(), 0)); - max = Math.max(max, NumericUtils.sortableBytesToLong(values.getMaxPackedValue(), 0)); - } - } - - if (min == Long.MAX_VALUE || max == Long.MIN_VALUE) { - return null; - } - return new long[] { min, max }; - } - - /** - * Finds the min and max bounds of the field for the segment - * - * @return null if the field is empty or not indexed - */ - private static long[] getSegmentBounds(final LeafReaderContext context, final String fieldName) throws IOException { - long min = Long.MAX_VALUE, max = Long.MIN_VALUE; - final PointValues values = context.reader().getPointValues(fieldName); - if (values != null) { - min = Math.min(min, NumericUtils.sortableBytesToLong(values.getMinPackedValue(), 0)); - max = Math.max(max, NumericUtils.sortableBytesToLong(values.getMaxPackedValue(), 0)); - } - - if (min == Long.MAX_VALUE || max == Long.MIN_VALUE) { - return null; - } - return new long[] { min, max }; - } - - /** - * Gets the min and max bounds of the field for the shard search - * Depending on the query part, the bounds are computed differently - * - * @return null if the processed query not supported by the optimization - */ - public static long[] getDateHistoAggBounds(final SearchContext context, final String fieldName) throws IOException { - final Query cq = unwrapIntoConcreteQuery(context.query()); - if (cq instanceof PointRangeQuery) { - final PointRangeQuery prq = (PointRangeQuery) cq; - final long[] indexBounds = getShardBounds(context, fieldName); - if (indexBounds == null) return null; - return getBoundsWithRangeQuery(prq, fieldName, indexBounds); - } else if (cq instanceof MatchAllDocsQuery) { - return getShardBounds(context, fieldName); - } else if (cq instanceof FieldExistsQuery) { - // when a range query covers all values of a shard, it will be rewrite field exists query - if (((FieldExistsQuery) cq).getField().equals(fieldName)) { - return getShardBounds(context, fieldName); - } - } - - return null; - } - - private static long[] getBoundsWithRangeQuery(PointRangeQuery prq, String fieldName, long[] indexBounds) { - // Ensure that the query and aggregation are on the same field - if (prq.getField().equals(fieldName)) { - // Minimum bound for aggregation is the max between query and global - long lower = Math.max(NumericUtils.sortableBytesToLong(prq.getLowerPoint(), 0), indexBounds[0]); - // Maximum bound for aggregation is the min between query and global - long upper = Math.min(NumericUtils.sortableBytesToLong(prq.getUpperPoint(), 0), indexBounds[1]); - if (lower > upper) { - return null; - } - return new long[] { lower, upper }; - } - - return null; - } - - /** - * Context object for fast filter optimization - *

- * Usage: first set aggregation type, then check isRewriteable, then buildFastFilter - */ - public static class FastFilterContext { - private boolean rewriteable = false; - private boolean rangesBuiltAtShardLevel = false; - - private AggregationType aggregationType; - private final SearchContext context; - - private MappedFieldType fieldType; - private Ranges ranges; - - // debug info related fields - public int leaf; - public int inner; - public int segments; - public int optimizedSegments; - - public FastFilterContext(SearchContext context) { - this.context = context; - } - - public FastFilterContext(SearchContext context, AggregationType aggregationType) { - this.context = context; - this.aggregationType = aggregationType; - } - - public AggregationType getAggregationType() { - return aggregationType; - } - - public void setAggregationType(AggregationType aggregationType) { - this.aggregationType = aggregationType; - } - - public boolean isRewriteable(final Object parent, final int subAggLength) { - if (context.maxAggRewriteFilters() == 0) return false; - - boolean rewriteable = aggregationType.isRewriteable(parent, subAggLength); - logger.debug("Fast filter rewriteable: {} for shard {}", rewriteable, context.indexShard().shardId()); - this.rewriteable = rewriteable; - return rewriteable; - } - - public void buildRanges(MappedFieldType fieldType) throws IOException { - assert ranges == null : "Ranges should only be built once at shard level, but they are already built"; - this.fieldType = fieldType; - this.ranges = this.aggregationType.buildRanges(context, fieldType); - if (ranges != null) { - logger.debug("Ranges built for shard {}", context.indexShard().shardId()); - rangesBuiltAtShardLevel = true; - } - } - - private Ranges buildRanges(LeafReaderContext leaf) throws IOException { - Ranges ranges = this.aggregationType.buildRanges(leaf, context, fieldType); - if (ranges != null) { - logger.debug("Ranges built for shard {} segment {}", context.indexShard().shardId(), leaf.ord); - } - return ranges; - } - - /** - * Try to populate the bucket doc counts for aggregation - *

- * Usage: invoked at segment level — in getLeafCollector of aggregator - * - * @param bucketOrd bucket ordinal producer - * @param incrementDocCount consume the doc_count results for certain ordinal - */ - public boolean tryFastFilterAggregation( - final LeafReaderContext ctx, - final BiConsumer incrementDocCount, - final Function bucketOrd - ) throws IOException { - this.segments++; - if (!this.rewriteable) { - return false; - } - - if (ctx.reader().hasDeletions()) return false; - - PointValues values = ctx.reader().getPointValues(this.fieldType.name()); - if (values == null) return false; - // only proceed if every document corresponds to exactly one point - if (values.getDocCount() != values.size()) return false; - - NumericDocValues docCountValues = DocValues.getNumeric(ctx.reader(), DocCountFieldMapper.NAME); - if (docCountValues.nextDoc() != NO_MORE_DOCS) { - logger.debug( - "Shard {} segment {} has at least one document with _doc_count field, skip fast filter optimization", - this.context.indexShard().shardId(), - ctx.ord - ); - return false; - } - - // even if no ranges built at shard level, we can still perform the optimization - // when functionally match-all at segment level - if (!this.rangesBuiltAtShardLevel && !segmentMatchAll(this.context, ctx)) { - return false; - } - - Ranges ranges = this.ranges; - if (ranges == null) { - logger.debug( - "Shard {} segment {} functionally match all documents. Build the fast filter", - this.context.indexShard().shardId(), - ctx.ord - ); - ranges = this.buildRanges(ctx); - if (ranges == null) { - return false; - } - } - - DebugInfo debugInfo = this.aggregationType.tryFastFilterAggregation(values, ranges, incrementDocCount, bucketOrd); - this.consumeDebugInfo(debugInfo); - - this.optimizedSegments++; - logger.debug("Fast filter optimization applied to shard {} segment {}", this.context.indexShard().shardId(), ctx.ord); - logger.debug("crossed leaf nodes: {}, inner nodes: {}", this.leaf, this.inner); - return true; - } - - private void consumeDebugInfo(DebugInfo debug) { - leaf += debug.leaf; - inner += debug.inner; - } - } - - /** - * Different types have different pre-conditions, filter building logic, etc. - */ - interface AggregationType { - boolean isRewriteable(Object parent, int subAggLength); - - Ranges buildRanges(SearchContext ctx, MappedFieldType fieldType) throws IOException; - - Ranges buildRanges(LeafReaderContext leaf, SearchContext ctx, MappedFieldType fieldType) throws IOException; - - DebugInfo tryFastFilterAggregation( - PointValues values, - Ranges ranges, - BiConsumer incrementDocCount, - Function bucketOrd - ) throws IOException; - } - - /** - * For date histogram aggregation - */ - public static abstract class AbstractDateHistogramAggregationType implements AggregationType { - private final MappedFieldType fieldType; - private final boolean missing; - private final boolean hasScript; - private LongBounds hardBounds; - - public AbstractDateHistogramAggregationType(MappedFieldType fieldType, boolean missing, boolean hasScript) { - this.fieldType = fieldType; - this.missing = missing; - this.hasScript = hasScript; - } - - public AbstractDateHistogramAggregationType(MappedFieldType fieldType, boolean missing, boolean hasScript, LongBounds hardBounds) { - this(fieldType, missing, hasScript); - this.hardBounds = hardBounds; - } - - @Override - public boolean isRewriteable(Object parent, int subAggLength) { - if (parent == null && subAggLength == 0 && !missing && !hasScript) { - if (fieldType != null && fieldType instanceof DateFieldMapper.DateFieldType) { - return fieldType.isSearchable(); - } - } - return false; - } - - @Override - public Ranges buildRanges(SearchContext context, MappedFieldType fieldType) throws IOException { - long[] bounds = getDateHistoAggBounds(context, fieldType.name()); - logger.debug("Bounds are {} for shard {}", bounds, context.indexShard().shardId()); - return buildRanges(context, bounds); - } - - @Override - public Ranges buildRanges(LeafReaderContext leaf, SearchContext context, MappedFieldType fieldType) throws IOException { - long[] bounds = getSegmentBounds(leaf, fieldType.name()); - logger.debug("Bounds are {} for shard {} segment {}", bounds, context.indexShard().shardId(), leaf.ord); - return buildRanges(context, bounds); - } - - private Ranges buildRanges(SearchContext context, long[] bounds) throws IOException { - bounds = processHardBounds(bounds); - if (bounds == null) { - return null; - } - assert bounds[0] <= bounds[1] : "Low bound should be less than high bound"; - - final Rounding rounding = getRounding(bounds[0], bounds[1]); - final OptionalLong intervalOpt = Rounding.getInterval(rounding); - if (intervalOpt.isEmpty()) { - return null; - } - final long interval = intervalOpt.getAsLong(); - - // process the after key of composite agg - processAfterKey(bounds, interval); - - return FastFilterRewriteHelper.createRangesFromAgg( - context, - (DateFieldMapper.DateFieldType) fieldType, - interval, - getRoundingPrepared(), - bounds[0], - bounds[1] - ); - } - - protected abstract Rounding getRounding(final long low, final long high); - - protected abstract Rounding.Prepared getRoundingPrepared(); - - protected void processAfterKey(long[] bound, long interval) {} - - protected long[] processHardBounds(long[] bounds) { - if (bounds != null) { - // Update min/max limit if user specified any hard bounds - if (hardBounds != null) { - if (hardBounds.getMin() > bounds[0]) { - bounds[0] = hardBounds.getMin(); - } - if (hardBounds.getMax() - 1 < bounds[1]) { - bounds[1] = hardBounds.getMax() - 1; // hard bounds max is exclusive - } - if (bounds[0] > bounds[1]) { - return null; - } - } - } - return bounds; - } - - public DateFieldMapper.DateFieldType getFieldType() { - assert fieldType instanceof DateFieldMapper.DateFieldType; - return (DateFieldMapper.DateFieldType) fieldType; - } - - @Override - public DebugInfo tryFastFilterAggregation( - PointValues values, - Ranges ranges, - BiConsumer incrementDocCount, - Function bucketOrd - ) throws IOException { - int size = Integer.MAX_VALUE; - if (this instanceof CompositeAggregator.CompositeAggregationType) { - size = ((CompositeAggregator.CompositeAggregationType) this).getSize(); - } - - DateFieldMapper.DateFieldType fieldType = getFieldType(); - BiConsumer incrementFunc = (activeIndex, docCount) -> { - long rangeStart = LongPoint.decodeDimension(ranges.lowers[activeIndex], 0); - rangeStart = fieldType.convertNanosToMillis(rangeStart); - long ord = getBucketOrd(bucketOrd.apply(rangeStart)); - incrementDocCount.accept(ord, (long) docCount); - }; - - return multiRangesTraverse(values.getPointTree(), ranges, incrementFunc, size); - } - - private static long getBucketOrd(long bucketOrd) { - if (bucketOrd < 0) { // already seen - bucketOrd = -1 - bucketOrd; - } - - return bucketOrd; - } - } - - /** - * For range aggregation - */ - public static class RangeAggregationType implements AggregationType { - - private final ValuesSourceConfig config; - private final Range[] ranges; - - public RangeAggregationType(ValuesSourceConfig config, Range[] ranges) { - this.config = config; - this.ranges = ranges; - } - - @Override - public boolean isRewriteable(Object parent, int subAggLength) { - if (config.fieldType() == null) return false; - MappedFieldType fieldType = config.fieldType(); - if (fieldType.isSearchable() == false || !(fieldType instanceof NumericPointEncoder)) return false; - - if (parent == null && subAggLength == 0 && config.script() == null && config.missing() == null) { - if (config.getValuesSource() instanceof ValuesSource.Numeric.FieldData) { - // ranges are already sorted by from and then to - // we want ranges not overlapping with each other - double prevTo = ranges[0].getTo(); - for (int i = 1; i < ranges.length; i++) { - if (prevTo > ranges[i].getFrom()) { - return false; - } - prevTo = ranges[i].getTo(); - } - return true; - } - } - return false; - } - - @Override - public Ranges buildRanges(SearchContext context, MappedFieldType fieldType) { - assert fieldType instanceof NumericPointEncoder; - NumericPointEncoder numericPointEncoder = (NumericPointEncoder) fieldType; - byte[][] lowers = new byte[ranges.length][]; - byte[][] uppers = new byte[ranges.length][]; - for (int i = 0; i < ranges.length; i++) { - double rangeMin = ranges[i].getFrom(); - double rangeMax = ranges[i].getTo(); - byte[] lower = numericPointEncoder.encodePoint(rangeMin); - byte[] upper = numericPointEncoder.encodePoint(rangeMax); - lowers[i] = lower; - uppers[i] = upper; - } - - return new Ranges(lowers, uppers); - } - - @Override - public Ranges buildRanges(LeafReaderContext leaf, SearchContext ctx, MappedFieldType fieldType) { - throw new UnsupportedOperationException("Range aggregation should not build ranges at segment level"); - } - - @Override - public DebugInfo tryFastFilterAggregation( - PointValues values, - Ranges ranges, - BiConsumer incrementDocCount, - Function bucketOrd - ) throws IOException { - int size = Integer.MAX_VALUE; - - BiConsumer incrementFunc = (activeIndex, docCount) -> { - long ord = bucketOrd.apply(activeIndex); - incrementDocCount.accept(ord, (long) docCount); - }; - - return multiRangesTraverse(values.getPointTree(), ranges, incrementFunc, size); - } - } - - public static boolean isCompositeAggRewriteable(CompositeValuesSourceConfig[] sourceConfigs) { - return sourceConfigs.length == 1 && sourceConfigs[0].valuesSource() instanceof RoundingValuesSource; - } - - private static boolean segmentMatchAll(SearchContext ctx, LeafReaderContext leafCtx) throws IOException { - Weight weight = ctx.searcher().createWeight(ctx.query(), ScoreMode.COMPLETE_NO_SCORES, 1f); - return weight != null && weight.count(leafCtx) == leafCtx.reader().numDocs(); - } - - /** - * Creates the date ranges from date histo aggregations using its interval, - * and min/max boundaries - */ - private static Ranges createRangesFromAgg( - final SearchContext context, - final DateFieldMapper.DateFieldType fieldType, - final long interval, - final Rounding.Prepared preparedRounding, - long low, - final long high - ) { - // Calculate the number of buckets using range and interval - long roundedLow = preparedRounding.round(fieldType.convertNanosToMillis(low)); - long prevRounded = roundedLow; - int bucketCount = 0; - while (roundedLow <= fieldType.convertNanosToMillis(high)) { - bucketCount++; - int maxNumFilterBuckets = context.maxAggRewriteFilters(); - if (bucketCount > maxNumFilterBuckets) { - logger.debug("Max number of filters reached [{}], skip the fast filter optimization", maxNumFilterBuckets); - return null; - } - // Below rounding is needed as the interval could return in - // non-rounded values for something like calendar month - roundedLow = preparedRounding.round(roundedLow + interval); - if (prevRounded == roundedLow) break; // prevents getting into an infinite loop - prevRounded = roundedLow; - } - - long[][] ranges = new long[bucketCount][2]; - if (bucketCount > 0) { - roundedLow = preparedRounding.round(fieldType.convertNanosToMillis(low)); - - int i = 0; - while (i < bucketCount) { - // Calculate the lower bucket bound - long lower = i == 0 ? low : fieldType.convertRoundedMillisToNanos(roundedLow); - roundedLow = preparedRounding.round(roundedLow + interval); - - // plus one on high value because upper bound is exclusive, but high value exists - long upper = i + 1 == bucketCount ? high + 1 : fieldType.convertRoundedMillisToNanos(roundedLow); - - ranges[i][0] = lower; - ranges[i][1] = upper; - i++; - } - } - - byte[][] lowers = new byte[ranges.length][]; - byte[][] uppers = new byte[ranges.length][]; - for (int i = 0; i < ranges.length; i++) { - byte[] lower = LONG.encodePoint(ranges[i][0]); - byte[] max = LONG.encodePoint(ranges[i][1]); - lowers[i] = lower; - uppers[i] = max; - } - - return new Ranges(lowers, uppers); - } - - /** - * @param maxNumNonZeroRanges the number of non-zero ranges to collect - */ - private static DebugInfo multiRangesTraverse( - final PointValues.PointTree tree, - final Ranges ranges, - final BiConsumer incrementDocCount, - final int maxNumNonZeroRanges - ) throws IOException { - DebugInfo debugInfo = new DebugInfo(); - int activeIndex = ranges.firstRangeIndex(tree.getMinPackedValue(), tree.getMaxPackedValue()); - if (activeIndex < 0) { - logger.debug("No ranges match the query, skip the fast filter optimization"); - return debugInfo; - } - RangeCollectorForPointTree collector = new RangeCollectorForPointTree(incrementDocCount, maxNumNonZeroRanges, ranges, activeIndex); - PointValues.IntersectVisitor visitor = getIntersectVisitor(collector); - try { - intersectWithRanges(visitor, tree, collector, debugInfo); - } catch (CollectionTerminatedException e) { - logger.debug("Early terminate since no more range to collect"); - } - collector.finalizePreviousRange(); - - return debugInfo; - } - - private static class Ranges { - byte[][] lowers; // inclusive - byte[][] uppers; // exclusive - int size; - int byteLen; - static ArrayUtil.ByteArrayComparator comparator; - - Ranges(byte[][] lowers, byte[][] uppers) { - this.lowers = lowers; - this.uppers = uppers; - assert lowers.length == uppers.length; - this.size = lowers.length; - this.byteLen = lowers[0].length; - comparator = ArrayUtil.getUnsignedComparator(byteLen); - } - - public int firstRangeIndex(byte[] globalMin, byte[] globalMax) { - if (compareByteValue(lowers[0], globalMax) > 0) { - return -1; - } - int i = 0; - while (compareByteValue(uppers[i], globalMin) <= 0) { - i++; - if (i >= size) { - return -1; - } - } - return i; - } - - public static int compareByteValue(byte[] value1, byte[] value2) { - return comparator.compare(value1, 0, value2, 0); - } - - public static boolean withinLowerBound(byte[] value, byte[] lowerBound) { - return compareByteValue(value, lowerBound) >= 0; - } - - public static boolean withinUpperBound(byte[] value, byte[] upperBound) { - return compareByteValue(value, upperBound) < 0; - } - } - - private static void intersectWithRanges( - PointValues.IntersectVisitor visitor, - PointValues.PointTree pointTree, - RangeCollectorForPointTree collector, - DebugInfo debug - ) throws IOException { - PointValues.Relation r = visitor.compare(pointTree.getMinPackedValue(), pointTree.getMaxPackedValue()); - - switch (r) { - case CELL_INSIDE_QUERY: - collector.countNode((int) pointTree.size()); - debug.visitInner(); - break; - case CELL_CROSSES_QUERY: - if (pointTree.moveToChild()) { - do { - intersectWithRanges(visitor, pointTree, collector, debug); - } while (pointTree.moveToSibling()); - pointTree.moveToParent(); - } else { - pointTree.visitDocValues(visitor); - debug.visitLeaf(); - } - break; - case CELL_OUTSIDE_QUERY: - } - } - - private static PointValues.IntersectVisitor getIntersectVisitor(RangeCollectorForPointTree collector) { - return new PointValues.IntersectVisitor() { - @Override - public void visit(int docID) throws IOException { - // this branch should be unreachable - throw new UnsupportedOperationException( - "This IntersectVisitor does not perform any actions on a " + "docID=" + docID + " node being visited" - ); - } - - @Override - public void visit(int docID, byte[] packedValue) throws IOException { - visitPoints(packedValue, collector::count); - } - - @Override - public void visit(DocIdSetIterator iterator, byte[] packedValue) throws IOException { - visitPoints(packedValue, () -> { - for (int doc = iterator.nextDoc(); doc != NO_MORE_DOCS; doc = iterator.nextDoc()) { - collector.count(); - } - }); - } - - private void visitPoints(byte[] packedValue, CheckedRunnable collect) throws IOException { - if (!collector.withinUpperBound(packedValue)) { - collector.finalizePreviousRange(); - if (collector.iterateRangeEnd(packedValue)) { - throw new CollectionTerminatedException(); - } - } - - if (collector.withinRange(packedValue)) { - collect.run(); - } - } - - @Override - public PointValues.Relation compare(byte[] minPackedValue, byte[] maxPackedValue) { - // try to find the first range that may collect values from this cell - if (!collector.withinUpperBound(minPackedValue)) { - collector.finalizePreviousRange(); - if (collector.iterateRangeEnd(minPackedValue)) { - throw new CollectionTerminatedException(); - } - } - // after the loop, min < upper - // cell could be outside [min max] lower - if (!collector.withinLowerBound(maxPackedValue)) { - return PointValues.Relation.CELL_OUTSIDE_QUERY; - } - if (collector.withinRange(minPackedValue) && collector.withinRange(maxPackedValue)) { - return PointValues.Relation.CELL_INSIDE_QUERY; - } - return PointValues.Relation.CELL_CROSSES_QUERY; - } - }; - } - - private static class RangeCollectorForPointTree { - private final BiConsumer incrementRangeDocCount; - private int counter = 0; - - private final Ranges ranges; - private int activeIndex; - - private int visitedRange = 0; - private final int maxNumNonZeroRange; - - public RangeCollectorForPointTree( - BiConsumer incrementRangeDocCount, - int maxNumNonZeroRange, - Ranges ranges, - int activeIndex - ) { - this.incrementRangeDocCount = incrementRangeDocCount; - this.maxNumNonZeroRange = maxNumNonZeroRange; - this.ranges = ranges; - this.activeIndex = activeIndex; - } - - private void count() { - counter++; - } - - private void countNode(int count) { - counter += count; - } - - private void finalizePreviousRange() { - if (counter > 0) { - incrementRangeDocCount.accept(activeIndex, counter); - counter = 0; - } - } - - /** - * @return true when iterator exhausted or collect enough non-zero ranges - */ - private boolean iterateRangeEnd(byte[] value) { - // the new value may not be contiguous to the previous one - // so try to find the first next range that cross the new value - while (!withinUpperBound(value)) { - if (++activeIndex >= ranges.size) { - return true; - } - } - visitedRange++; - return visitedRange > maxNumNonZeroRange; - } - - private boolean withinLowerBound(byte[] value) { - return Ranges.withinLowerBound(value, ranges.lowers[activeIndex]); - } - - private boolean withinUpperBound(byte[] value) { - return Ranges.withinUpperBound(value, ranges.uppers[activeIndex]); - } - - private boolean withinRange(byte[] value) { - return withinLowerBound(value) && withinUpperBound(value); - } - } - - /** - * Contains debug info of BKD traversal to show in profile - */ - private static class DebugInfo { - private int leaf = 0; // leaf node visited - private int inner = 0; // inner node visited - - private void visitLeaf() { - leaf++; - } - - private void visitInner() { - inner++; - } - } -} diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/composite/CompositeAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/composite/CompositeAggregator.java index bfb484dcf478d..cfe716eb57ca8 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/bucket/composite/CompositeAggregator.java +++ b/server/src/main/java/org/opensearch/search/aggregations/bucket/composite/CompositeAggregator.java @@ -73,8 +73,8 @@ import org.opensearch.search.aggregations.MultiBucketCollector; import org.opensearch.search.aggregations.MultiBucketConsumerService; import org.opensearch.search.aggregations.bucket.BucketsAggregator; -import org.opensearch.search.aggregations.bucket.FastFilterRewriteHelper; -import org.opensearch.search.aggregations.bucket.FastFilterRewriteHelper.AbstractDateHistogramAggregationType; +import org.opensearch.search.aggregations.bucket.filterrewrite.CompositeAggregatorBridge; +import org.opensearch.search.aggregations.bucket.filterrewrite.FilterRewriteOptimizationContext; import org.opensearch.search.aggregations.bucket.missing.MissingOrder; import org.opensearch.search.aggregations.bucket.terms.LongKeyedBucketOrds; import org.opensearch.search.internal.SearchContext; @@ -89,13 +89,15 @@ import java.util.List; import java.util.Map; import java.util.function.BiConsumer; +import java.util.function.Function; import java.util.function.LongUnaryOperator; import java.util.stream.Collectors; import static org.opensearch.search.aggregations.MultiBucketConsumerService.MAX_BUCKET_SETTING; +import static org.opensearch.search.aggregations.bucket.filterrewrite.DateHistogramAggregatorBridge.segmentMatchAll; /** - * Main aggregator that aggregates docs from mulitple aggregations + * Main aggregator that aggregates docs from multiple aggregations * * @opensearch.internal */ @@ -118,9 +120,8 @@ public final class CompositeAggregator extends BucketsAggregator { private boolean earlyTerminated; - private final FastFilterRewriteHelper.FastFilterContext fastFilterContext; - private LongKeyedBucketOrds bucketOrds = null; - private Rounding.Prepared preparedRounding = null; + private final FilterRewriteOptimizationContext filterRewriteOptimizationContext; + private LongKeyedBucketOrds bucketOrds; CompositeAggregator( String name, @@ -166,57 +167,62 @@ public final class CompositeAggregator extends BucketsAggregator { this.queue = new CompositeValuesCollectorQueue(context.bigArrays(), sources, size, rawAfterKey); this.rawAfterKey = rawAfterKey; - fastFilterContext = new FastFilterRewriteHelper.FastFilterContext(context); - if (!FastFilterRewriteHelper.isCompositeAggRewriteable(sourceConfigs)) { - return; - } - fastFilterContext.setAggregationType(new CompositeAggregationType()); - if (fastFilterContext.isRewriteable(parent, subAggregators.length)) { - // bucketOrds is used for saving date histogram results - bucketOrds = LongKeyedBucketOrds.build(context.bigArrays(), CardinalityUpperBound.ONE); - preparedRounding = ((CompositeAggregationType) fastFilterContext.getAggregationType()).getRoundingPrepared(); - fastFilterContext.buildRanges(sourceConfigs[0].fieldType()); - } - } + CompositeAggregatorBridge bridge = new CompositeAggregatorBridge() { + private RoundingValuesSource valuesSource; + private long afterKey = -1L; - /** - * Currently the filter rewrite is only supported for date histograms - */ - public class CompositeAggregationType extends AbstractDateHistogramAggregationType { - private final RoundingValuesSource valuesSource; - private long afterKey = -1L; - - public CompositeAggregationType() { - super(sourceConfigs[0].fieldType(), sourceConfigs[0].missingBucket(), sourceConfigs[0].hasScript()); - this.valuesSource = (RoundingValuesSource) sourceConfigs[0].valuesSource(); - if (rawAfterKey != null) { - assert rawAfterKey.size() == 1 && formats.size() == 1; - this.afterKey = formats.get(0).parseLong(rawAfterKey.get(0).toString(), false, () -> { - throw new IllegalArgumentException("now() is not supported in [after] key"); - }); + @Override + protected boolean canOptimize() { + if (canOptimize(sourceConfigs)) { + this.valuesSource = (RoundingValuesSource) sourceConfigs[0].valuesSource(); + if (rawAfterKey != null) { + assert rawAfterKey.size() == 1 && formats.size() == 1; + this.afterKey = formats.get(0).parseLong(rawAfterKey.get(0).toString(), false, () -> { + throw new IllegalArgumentException("now() is not supported in [after] key"); + }); + } + + // bucketOrds is used for saving the date histogram results got from the optimization path + bucketOrds = LongKeyedBucketOrds.build(context.bigArrays(), CardinalityUpperBound.ONE); + return true; + } + return false; } - } - public Rounding getRounding(final long low, final long high) { - return valuesSource.getRounding(); - } + @Override + protected void prepare() throws IOException { + buildRanges(context); + } - public Rounding.Prepared getRoundingPrepared() { - return valuesSource.getPreparedRounding(); - } + protected Rounding getRounding(final long low, final long high) { + return valuesSource.getRounding(); + } - @Override - protected void processAfterKey(long[] bound, long interval) { - // afterKey is the last bucket key in previous response, and the bucket key - // is the minimum of all values in the bucket, so need to add the interval - if (afterKey != -1L) { - bound[0] = afterKey + interval; + protected Rounding.Prepared getRoundingPrepared() { + return valuesSource.getPreparedRounding(); } - } - public int getSize() { - return size; - } + @Override + protected long[] processAfterKey(long[] bounds, long interval) { + // afterKey is the last bucket key in previous response, and the bucket key + // is the minimum of all values in the bucket, so need to add the interval + if (afterKey != -1L) { + bounds[0] = afterKey + interval; + } + return bounds; + } + + @Override + protected int getSize() { + return size; + } + + @Override + protected Function bucketOrdProducer() { + return (key) -> bucketOrds.add(0, getRoundingPrepared().round((long) key)); + } + }; + filterRewriteOptimizationContext = new FilterRewriteOptimizationContext(bridge, parent, subAggregators.length, context); } @Override @@ -368,7 +374,7 @@ private boolean isMaybeMultivalued(LeafReaderContext context, SortField sortFiel return v2 != null && DocValues.unwrapSingleton(v2) == null; default: - // we have no clue whether the field is multi-valued or not so we assume it is. + // we have no clue whether the field is multivalued or not so we assume it is. return true; } } @@ -551,11 +557,7 @@ private void processLeafFromQuery(LeafReaderContext ctx, Sort indexSortPrefix) t @Override protected LeafBucketCollector getLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub) throws IOException { - boolean optimized = fastFilterContext.tryFastFilterAggregation( - ctx, - this::incrementBucketDocCount, - (key) -> bucketOrds.add(0, preparedRounding.round((long) key)) - ); + boolean optimized = filterRewriteOptimizationContext.tryOptimize(ctx, this::incrementBucketDocCount, segmentMatchAll(context, ctx)); if (optimized) throw new CollectionTerminatedException(); finishLeaf(); @@ -709,11 +711,6 @@ private static class Entry { @Override public void collectDebugInfo(BiConsumer add) { - if (fastFilterContext.optimizedSegments > 0) { - add.accept("optimized_segments", fastFilterContext.optimizedSegments); - add.accept("unoptimized_segments", fastFilterContext.segments - fastFilterContext.optimizedSegments); - add.accept("leaf_visited", fastFilterContext.leaf); - add.accept("inner_visited", fastFilterContext.inner); - } + filterRewriteOptimizationContext.populateDebugInfo(add); } } diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/AggregatorBridge.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/AggregatorBridge.java new file mode 100644 index 0000000000000..6b34582b259ea --- /dev/null +++ b/server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/AggregatorBridge.java @@ -0,0 +1,84 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.search.aggregations.bucket.filterrewrite; + +import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.index.PointValues; +import org.opensearch.index.mapper.MappedFieldType; + +import java.io.IOException; +import java.util.function.BiConsumer; +import java.util.function.Consumer; + +/** + * This interface provides a bridge between an aggregator and the optimization context, allowing + * the aggregator to provide data and optimize the aggregation process. + * + *

The main purpose of this interface is to encapsulate the aggregator-specific optimization + * logic and provide access to the data in Aggregator that is required for optimization, while keeping the optimization + * business logic separate from the aggregator implementation. + * + *

To use this interface to optimize an aggregator, you should subclass this interface in this package + * and put any specific optimization business logic in it. Then implement this subclass in the aggregator + * to provide data that is needed for doing the optimization + * + * @opensearch.internal + */ +public abstract class AggregatorBridge { + + /** + * The field type associated with this aggregator bridge. + */ + MappedFieldType fieldType; + + Consumer setRanges; + + void setRangesConsumer(Consumer setRanges) { + this.setRanges = setRanges; + } + + /** + * Checks whether the aggregator can be optimized. + *

+ * This method is supposed to be implemented in a specific aggregator to take in fields from there + * + * @return {@code true} if the aggregator can be optimized, {@code false} otherwise. + * The result will be saved in the optimization context. + */ + protected abstract boolean canOptimize(); + + /** + * Prepares the optimization at shard level after checking aggregator is optimizable. + *

+ * For example, figure out what are the ranges from the aggregation to do the optimization later + *

+ * This method is supposed to be implemented in a specific aggregator to take in fields from there + */ + protected abstract void prepare() throws IOException; + + /** + * Prepares the optimization for a specific segment when the segment is functionally matching all docs + * + * @param leaf the leaf reader context for the segment + */ + abstract Ranges tryBuildRangesFromSegment(LeafReaderContext leaf) throws IOException; + + /** + * Attempts to build aggregation results for a segment + * + * @param values the point values (index structure for numeric values) for a segment + * @param incrementDocCount a consumer to increment the document count for a range bucket. The First parameter is document count, the second is the key of the bucket + * @param ranges + */ + abstract FilterRewriteOptimizationContext.DebugInfo tryOptimize( + PointValues values, + BiConsumer incrementDocCount, + Ranges ranges + ) throws IOException; +} diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/CompositeAggregatorBridge.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/CompositeAggregatorBridge.java new file mode 100644 index 0000000000000..e122e7bda0b6a --- /dev/null +++ b/server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/CompositeAggregatorBridge.java @@ -0,0 +1,36 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.search.aggregations.bucket.filterrewrite; + +import org.opensearch.index.mapper.DateFieldMapper; +import org.opensearch.index.mapper.MappedFieldType; +import org.opensearch.search.aggregations.bucket.composite.CompositeValuesSourceConfig; +import org.opensearch.search.aggregations.bucket.composite.RoundingValuesSource; + +/** + * For composite aggregation to do optimization when it only has a single date histogram source + */ +public abstract class CompositeAggregatorBridge extends DateHistogramAggregatorBridge { + protected boolean canOptimize(CompositeValuesSourceConfig[] sourceConfigs) { + if (sourceConfigs.length != 1 || !(sourceConfigs[0].valuesSource() instanceof RoundingValuesSource)) return false; + return canOptimize(sourceConfigs[0].missingBucket(), sourceConfigs[0].hasScript(), sourceConfigs[0].fieldType()); + } + + private boolean canOptimize(boolean missing, boolean hasScript, MappedFieldType fieldType) { + if (!missing && !hasScript) { + if (fieldType instanceof DateFieldMapper.DateFieldType) { + if (fieldType.isSearchable()) { + this.fieldType = fieldType; + return true; + } + } + } + return false; + } +} diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/DateHistogramAggregatorBridge.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/DateHistogramAggregatorBridge.java new file mode 100644 index 0000000000000..8bff3fdc5fefb --- /dev/null +++ b/server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/DateHistogramAggregatorBridge.java @@ -0,0 +1,174 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.search.aggregations.bucket.filterrewrite; + +import org.apache.lucene.document.LongPoint; +import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.index.PointValues; +import org.apache.lucene.search.ScoreMode; +import org.apache.lucene.search.Weight; +import org.opensearch.common.Rounding; +import org.opensearch.index.mapper.DateFieldMapper; +import org.opensearch.index.mapper.MappedFieldType; +import org.opensearch.search.aggregations.bucket.histogram.LongBounds; +import org.opensearch.search.aggregations.support.ValuesSourceConfig; +import org.opensearch.search.internal.SearchContext; + +import java.io.IOException; +import java.util.OptionalLong; +import java.util.function.BiConsumer; +import java.util.function.Function; + +import static org.opensearch.search.aggregations.bucket.filterrewrite.PointTreeTraversal.multiRangesTraverse; + +/** + * For date histogram aggregation + */ +public abstract class DateHistogramAggregatorBridge extends AggregatorBridge { + + int maxRewriteFilters; + + protected boolean canOptimize(ValuesSourceConfig config) { + if (config.script() == null && config.missing() == null) { + MappedFieldType fieldType = config.fieldType(); + if (fieldType instanceof DateFieldMapper.DateFieldType) { + if (fieldType.isSearchable()) { + this.fieldType = fieldType; + return true; + } + } + } + return false; + } + + protected void buildRanges(SearchContext context) throws IOException { + long[] bounds = Helper.getDateHistoAggBounds(context, fieldType.name()); + this.maxRewriteFilters = context.maxAggRewriteFilters(); + setRanges.accept(buildRanges(bounds, maxRewriteFilters)); + } + + @Override + final Ranges tryBuildRangesFromSegment(LeafReaderContext leaf) throws IOException { + long[] bounds = Helper.getSegmentBounds(leaf, fieldType.name()); + return buildRanges(bounds, maxRewriteFilters); + } + + private Ranges buildRanges(long[] bounds, int maxRewriteFilters) { + bounds = processHardBounds(bounds); + if (bounds == null) { + return null; + } + assert bounds[0] <= bounds[1] : "Low bound should be less than high bound"; + + final Rounding rounding = getRounding(bounds[0], bounds[1]); + final OptionalLong intervalOpt = Rounding.getInterval(rounding); + if (intervalOpt.isEmpty()) { + return null; + } + final long interval = intervalOpt.getAsLong(); + + // process the after key of composite agg + bounds = processAfterKey(bounds, interval); + + return Helper.createRangesFromAgg( + (DateFieldMapper.DateFieldType) fieldType, + interval, + getRoundingPrepared(), + bounds[0], + bounds[1], + maxRewriteFilters + ); + } + + protected abstract Rounding getRounding(final long low, final long high); + + protected abstract Rounding.Prepared getRoundingPrepared(); + + protected long[] processAfterKey(long[] bounds, long interval) { + return bounds; + } + + protected long[] processHardBounds(long[] bounds) { + return processHardBounds(bounds, null); + } + + protected long[] processHardBounds(long[] bounds, LongBounds hardBounds) { + if (bounds != null) { + // Update min/max limit if user specified any hard bounds + if (hardBounds != null) { + if (hardBounds.getMin() > bounds[0]) { + bounds[0] = hardBounds.getMin(); + } + if (hardBounds.getMax() - 1 < bounds[1]) { + bounds[1] = hardBounds.getMax() - 1; // hard bounds max is exclusive + } + if (bounds[0] > bounds[1]) { + return null; + } + } + } + return bounds; + } + + private DateFieldMapper.DateFieldType getFieldType() { + assert fieldType instanceof DateFieldMapper.DateFieldType; + return (DateFieldMapper.DateFieldType) fieldType; + } + + protected int getSize() { + return Integer.MAX_VALUE; + } + + @Override + final FilterRewriteOptimizationContext.DebugInfo tryOptimize( + PointValues values, + BiConsumer incrementDocCount, + Ranges ranges + ) throws IOException { + int size = getSize(); + + DateFieldMapper.DateFieldType fieldType = getFieldType(); + BiConsumer incrementFunc = (activeIndex, docCount) -> { + long rangeStart = LongPoint.decodeDimension(ranges.lowers[activeIndex], 0); + rangeStart = fieldType.convertNanosToMillis(rangeStart); + long bucketOrd = getBucketOrd(bucketOrdProducer().apply(rangeStart)); + incrementDocCount.accept(bucketOrd, (long) docCount); + }; + + return multiRangesTraverse(values.getPointTree(), ranges, incrementFunc, size); + } + + private static long getBucketOrd(long bucketOrd) { + if (bucketOrd < 0) { // already seen + bucketOrd = -1 - bucketOrd; + } + + return bucketOrd; + } + + /** + * Provides a function to produce bucket ordinals from the lower bound of the range + */ + protected abstract Function bucketOrdProducer(); + + /** + * Checks whether the top level query matches all documents on the segment + * + *

This method creates a weight from the search context's query and checks whether the weight's + * document count matches the total number of documents in the leaf reader context. + * + * @param ctx the search context + * @param leafCtx the leaf reader context for the segment + * @return {@code true} if the segment matches all documents, {@code false} otherwise + */ + public static boolean segmentMatchAll(SearchContext ctx, LeafReaderContext leafCtx) throws IOException { + Weight weight = ctx.query().rewrite(ctx.searcher()).createWeight(ctx.searcher(), ScoreMode.COMPLETE_NO_SCORES, 1f); + return weight != null && weight.count(leafCtx) == leafCtx.reader().numDocs(); + } +} diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/FilterRewriteOptimizationContext.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/FilterRewriteOptimizationContext.java new file mode 100644 index 0000000000000..87faafe4526de --- /dev/null +++ b/server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/FilterRewriteOptimizationContext.java @@ -0,0 +1,189 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.search.aggregations.bucket.filterrewrite; + +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.apache.lucene.index.DocValues; +import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.index.NumericDocValues; +import org.apache.lucene.index.PointValues; +import org.opensearch.index.mapper.DocCountFieldMapper; +import org.opensearch.search.internal.SearchContext; + +import java.io.IOException; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.BiConsumer; + +import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; + +/** + * Context object for doing the filter rewrite optimization in ranges type aggregation + *

+ * This holds the common business logic and delegate aggregator-specific logic to {@link AggregatorBridge} + * + * @opensearch.internal + */ +public final class FilterRewriteOptimizationContext { + + private static final Logger logger = LogManager.getLogger(Helper.loggerName); + + private final boolean canOptimize; + private boolean preparedAtShardLevel = false; + + private final AggregatorBridge aggregatorBridge; + private String shardId; + + private Ranges ranges; // built at shard level + + // debug info related fields + private final AtomicInteger leafNodeVisited = new AtomicInteger(); + private final AtomicInteger innerNodeVisited = new AtomicInteger(); + private final AtomicInteger segments = new AtomicInteger(); + private final AtomicInteger optimizedSegments = new AtomicInteger(); + + public FilterRewriteOptimizationContext( + AggregatorBridge aggregatorBridge, + final Object parent, + final int subAggLength, + SearchContext context + ) throws IOException { + this.aggregatorBridge = aggregatorBridge; + this.canOptimize = this.canOptimize(parent, subAggLength, context); + } + + /** + * common logic for checking whether the optimization can be applied and prepare at shard level + * if the aggregation has any special logic, it should be done using {@link AggregatorBridge} + */ + private boolean canOptimize(final Object parent, final int subAggLength, SearchContext context) throws IOException { + if (context.maxAggRewriteFilters() == 0) return false; + + if (parent != null || subAggLength != 0) return false; + + boolean canOptimize = aggregatorBridge.canOptimize(); + if (canOptimize) { + aggregatorBridge.setRangesConsumer(this::setRanges); + + this.shardId = context.indexShard().shardId().toString(); + + assert ranges == null : "Ranges should only be built once at shard level, but they are already built"; + aggregatorBridge.prepare(); + if (ranges != null) { + preparedAtShardLevel = true; + } + } + logger.debug("Fast filter rewriteable: {} for shard {}", canOptimize, shardId); + + return canOptimize; + } + + void setRanges(Ranges ranges) { + this.ranges = ranges; + } + + /** + * Try to populate the bucket doc counts for aggregation + *

+ * Usage: invoked at segment level — in getLeafCollector of aggregator + * + * @param incrementDocCount consume the doc_count results for certain ordinal + * @param segmentMatchAll if your optimization can prepareFromSegment, you should pass in this flag to decide whether to prepareFromSegment + */ + public boolean tryOptimize(final LeafReaderContext leafCtx, final BiConsumer incrementDocCount, boolean segmentMatchAll) + throws IOException { + segments.incrementAndGet(); + if (!canOptimize) { + return false; + } + + if (leafCtx.reader().hasDeletions()) return false; + + PointValues values = leafCtx.reader().getPointValues(aggregatorBridge.fieldType.name()); + if (values == null) return false; + // only proceed if every document corresponds to exactly one point + if (values.getDocCount() != values.size()) return false; + + NumericDocValues docCountValues = DocValues.getNumeric(leafCtx.reader(), DocCountFieldMapper.NAME); + if (docCountValues.nextDoc() != NO_MORE_DOCS) { + logger.debug( + "Shard {} segment {} has at least one document with _doc_count field, skip fast filter optimization", + shardId, + leafCtx.ord + ); + return false; + } + + Ranges ranges = getRanges(leafCtx, segmentMatchAll); + if (ranges == null) return false; + + consumeDebugInfo(aggregatorBridge.tryOptimize(values, incrementDocCount, ranges)); + + optimizedSegments.incrementAndGet(); + logger.debug("Fast filter optimization applied to shard {} segment {}", shardId, leafCtx.ord); + logger.debug("Crossed leaf nodes: {}, inner nodes: {}", leafNodeVisited, innerNodeVisited); + + return true; + } + + Ranges getRanges(LeafReaderContext leafCtx, boolean segmentMatchAll) { + if (!preparedAtShardLevel) { + try { + return getRangesFromSegment(leafCtx, segmentMatchAll); + } catch (IOException e) { + logger.warn("Failed to build ranges from segment.", e); + return null; + } + } + return ranges; + } + + /** + * Even when ranges cannot be built at shard level, we can still build ranges + * at segment level when it's functionally match-all at segment level + */ + private Ranges getRangesFromSegment(LeafReaderContext leafCtx, boolean segmentMatchAll) throws IOException { + if (!segmentMatchAll) { + return null; + } + + logger.debug("Shard {} segment {} functionally match all documents. Build the fast filter", shardId, leafCtx.ord); + return aggregatorBridge.tryBuildRangesFromSegment(leafCtx); + } + + /** + * Contains debug info of BKD traversal to show in profile + */ + static class DebugInfo { + private final AtomicInteger leafNodeVisited = new AtomicInteger(); // leaf node visited + private final AtomicInteger innerNodeVisited = new AtomicInteger(); // inner node visited + + void visitLeaf() { + leafNodeVisited.incrementAndGet(); + } + + void visitInner() { + innerNodeVisited.incrementAndGet(); + } + } + + void consumeDebugInfo(DebugInfo debug) { + leafNodeVisited.addAndGet(debug.leafNodeVisited.get()); + innerNodeVisited.addAndGet(debug.innerNodeVisited.get()); + } + + public void populateDebugInfo(BiConsumer add) { + if (optimizedSegments.get() > 0) { + add.accept("optimized_segments", optimizedSegments.get()); + add.accept("unoptimized_segments", segments.get() - optimizedSegments.get()); + add.accept("leaf_visited", leafNodeVisited.get()); + add.accept("inner_visited", innerNodeVisited.get()); + } + } +} diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/Helper.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/Helper.java new file mode 100644 index 0000000000000..7493754d8efa2 --- /dev/null +++ b/server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/Helper.java @@ -0,0 +1,213 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.search.aggregations.bucket.filterrewrite; + +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.index.PointValues; +import org.apache.lucene.search.ConstantScoreQuery; +import org.apache.lucene.search.FieldExistsQuery; +import org.apache.lucene.search.IndexOrDocValuesQuery; +import org.apache.lucene.search.MatchAllDocsQuery; +import org.apache.lucene.search.PointRangeQuery; +import org.apache.lucene.search.Query; +import org.apache.lucene.util.NumericUtils; +import org.opensearch.common.Rounding; +import org.opensearch.common.lucene.search.function.FunctionScoreQuery; +import org.opensearch.index.mapper.DateFieldMapper; +import org.opensearch.index.query.DateRangeIncludingNowQuery; +import org.opensearch.search.internal.SearchContext; + +import java.io.IOException; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.function.Function; + +import static org.opensearch.index.mapper.NumberFieldMapper.NumberType.LONG; + +/** + * Utility class to help range filters rewrite optimization + * + * @opensearch.internal + */ +final class Helper { + + private Helper() {} + + static final String loggerName = Helper.class.getPackageName(); + private static final Logger logger = LogManager.getLogger(loggerName); + + private static final Map, Function> queryWrappers; + + // Initialize the wrapper map for unwrapping the query + static { + queryWrappers = new HashMap<>(); + queryWrappers.put(ConstantScoreQuery.class, q -> ((ConstantScoreQuery) q).getQuery()); + queryWrappers.put(FunctionScoreQuery.class, q -> ((FunctionScoreQuery) q).getSubQuery()); + queryWrappers.put(DateRangeIncludingNowQuery.class, q -> ((DateRangeIncludingNowQuery) q).getQuery()); + queryWrappers.put(IndexOrDocValuesQuery.class, q -> ((IndexOrDocValuesQuery) q).getIndexQuery()); + } + + /** + * Recursively unwraps query into the concrete form + * for applying the optimization + */ + private static Query unwrapIntoConcreteQuery(Query query) { + while (queryWrappers.containsKey(query.getClass())) { + query = queryWrappers.get(query.getClass()).apply(query); + } + + return query; + } + + /** + * Finds the global min and max bounds of the field for the shard across all segments + * + * @return null if the field is empty or not indexed + */ + private static long[] getShardBounds(final List leaves, final String fieldName) throws IOException { + long min = Long.MAX_VALUE, max = Long.MIN_VALUE; + for (LeafReaderContext leaf : leaves) { + final PointValues values = leaf.reader().getPointValues(fieldName); + if (values != null) { + min = Math.min(min, NumericUtils.sortableBytesToLong(values.getMinPackedValue(), 0)); + max = Math.max(max, NumericUtils.sortableBytesToLong(values.getMaxPackedValue(), 0)); + } + } + + if (min == Long.MAX_VALUE || max == Long.MIN_VALUE) { + return null; + } + return new long[] { min, max }; + } + + /** + * Finds the min and max bounds of the field for the segment + * + * @return null if the field is empty or not indexed + */ + static long[] getSegmentBounds(final LeafReaderContext context, final String fieldName) throws IOException { + long min = Long.MAX_VALUE, max = Long.MIN_VALUE; + final PointValues values = context.reader().getPointValues(fieldName); + if (values != null) { + min = Math.min(min, NumericUtils.sortableBytesToLong(values.getMinPackedValue(), 0)); + max = Math.max(max, NumericUtils.sortableBytesToLong(values.getMaxPackedValue(), 0)); + } + + if (min == Long.MAX_VALUE || max == Long.MIN_VALUE) { + return null; + } + return new long[] { min, max }; + } + + /** + * Gets the min and max bounds of the field for the shard search + * Depending on the query part, the bounds are computed differently + * + * @return null if the processed query not supported by the optimization + */ + public static long[] getDateHistoAggBounds(final SearchContext context, final String fieldName) throws IOException { + final Query cq = unwrapIntoConcreteQuery(context.query()); + final List leaves = context.searcher().getIndexReader().leaves(); + + if (cq instanceof PointRangeQuery) { + final PointRangeQuery prq = (PointRangeQuery) cq; + final long[] indexBounds = getShardBounds(leaves, fieldName); + if (indexBounds == null) return null; + return getBoundsWithRangeQuery(prq, fieldName, indexBounds); + } else if (cq instanceof MatchAllDocsQuery) { + return getShardBounds(leaves, fieldName); + } else if (cq instanceof FieldExistsQuery) { + // when a range query covers all values of a shard, it will be rewrite field exists query + if (((FieldExistsQuery) cq).getField().equals(fieldName)) { + return getShardBounds(leaves, fieldName); + } + } + + return null; + } + + private static long[] getBoundsWithRangeQuery(PointRangeQuery prq, String fieldName, long[] indexBounds) { + // Ensure that the query and aggregation are on the same field + if (prq.getField().equals(fieldName)) { + // Minimum bound for aggregation is the max between query and global + long lower = Math.max(NumericUtils.sortableBytesToLong(prq.getLowerPoint(), 0), indexBounds[0]); + // Maximum bound for aggregation is the min between query and global + long upper = Math.min(NumericUtils.sortableBytesToLong(prq.getUpperPoint(), 0), indexBounds[1]); + if (lower > upper) { + return null; + } + return new long[] { lower, upper }; + } + + return null; + } + + /** + * Creates the date ranges from date histo aggregations using its interval, + * and min/max boundaries + */ + static Ranges createRangesFromAgg( + final DateFieldMapper.DateFieldType fieldType, + final long interval, + final Rounding.Prepared preparedRounding, + long low, + final long high, + final int maxAggRewriteFilters + ) { + // Calculate the number of buckets using range and interval + long roundedLow = preparedRounding.round(fieldType.convertNanosToMillis(low)); + long prevRounded = roundedLow; + int bucketCount = 0; + while (roundedLow <= fieldType.convertNanosToMillis(high)) { + bucketCount++; + if (bucketCount > maxAggRewriteFilters) { + logger.debug("Max number of range filters reached [{}], skip the optimization", maxAggRewriteFilters); + return null; + } + // Below rounding is needed as the interval could return in + // non-rounded values for something like calendar month + roundedLow = preparedRounding.round(roundedLow + interval); + if (prevRounded == roundedLow) break; // prevents getting into an infinite loop + prevRounded = roundedLow; + } + + long[][] ranges = new long[bucketCount][2]; + if (bucketCount > 0) { + roundedLow = preparedRounding.round(fieldType.convertNanosToMillis(low)); + + int i = 0; + while (i < bucketCount) { + // Calculate the lower bucket bound + long lower = i == 0 ? low : fieldType.convertRoundedMillisToNanos(roundedLow); + roundedLow = preparedRounding.round(roundedLow + interval); + + // plus one on high value because upper bound is exclusive, but high value exists + long upper = i + 1 == bucketCount ? high + 1 : fieldType.convertRoundedMillisToNanos(roundedLow); + + ranges[i][0] = lower; + ranges[i][1] = upper; + i++; + } + } + + byte[][] lowers = new byte[ranges.length][]; + byte[][] uppers = new byte[ranges.length][]; + for (int i = 0; i < ranges.length; i++) { + byte[] lower = LONG.encodePoint(ranges[i][0]); + byte[] max = LONG.encodePoint(ranges[i][1]); + lowers[i] = lower; + uppers[i] = max; + } + + return new Ranges(lowers, uppers); + } +} diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/PointTreeTraversal.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/PointTreeTraversal.java new file mode 100644 index 0000000000000..581ecc416f486 --- /dev/null +++ b/server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/PointTreeTraversal.java @@ -0,0 +1,223 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.search.aggregations.bucket.filterrewrite; + +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.apache.lucene.index.PointValues; +import org.apache.lucene.search.CollectionTerminatedException; +import org.apache.lucene.search.DocIdSetIterator; +import org.opensearch.common.CheckedRunnable; + +import java.io.IOException; +import java.util.function.BiConsumer; + +import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; + +/** + * Utility class for traversing a {@link PointValues.PointTree} and collecting document counts for the ranges. + * + *

The main entry point is the {@link #multiRangesTraverse(PointValues.PointTree, Ranges, + * BiConsumer, int)} method + * + *

The class uses a {@link RangeCollectorForPointTree} to keep track of the active ranges and + * determine which parts of the tree to visit. The {@link + * PointValues.IntersectVisitor} implementation is responsible for the actual visitation and + * document count collection. + */ +final class PointTreeTraversal { + private PointTreeTraversal() {} + + private static final Logger logger = LogManager.getLogger(Helper.loggerName); + + /** + * Traverses the given {@link PointValues.PointTree} and collects document counts for the intersecting ranges. + * + * @param tree the point tree to traverse + * @param ranges the set of ranges to intersect with + * @param incrementDocCount a callback to increment the document count for a range bucket + * @param maxNumNonZeroRanges the maximum number of non-zero ranges to collect + * @return a {@link FilterRewriteOptimizationContext.DebugInfo} object containing debug information about the traversal + */ + static FilterRewriteOptimizationContext.DebugInfo multiRangesTraverse( + final PointValues.PointTree tree, + final Ranges ranges, + final BiConsumer incrementDocCount, + final int maxNumNonZeroRanges + ) throws IOException { + FilterRewriteOptimizationContext.DebugInfo debugInfo = new FilterRewriteOptimizationContext.DebugInfo(); + int activeIndex = ranges.firstRangeIndex(tree.getMinPackedValue(), tree.getMaxPackedValue()); + if (activeIndex < 0) { + logger.debug("No ranges match the query, skip the fast filter optimization"); + return debugInfo; + } + RangeCollectorForPointTree collector = new RangeCollectorForPointTree(incrementDocCount, maxNumNonZeroRanges, ranges, activeIndex); + PointValues.IntersectVisitor visitor = getIntersectVisitor(collector); + try { + intersectWithRanges(visitor, tree, collector, debugInfo); + } catch (CollectionTerminatedException e) { + logger.debug("Early terminate since no more range to collect"); + } + collector.finalizePreviousRange(); + + return debugInfo; + } + + private static void intersectWithRanges( + PointValues.IntersectVisitor visitor, + PointValues.PointTree pointTree, + RangeCollectorForPointTree collector, + FilterRewriteOptimizationContext.DebugInfo debug + ) throws IOException { + PointValues.Relation r = visitor.compare(pointTree.getMinPackedValue(), pointTree.getMaxPackedValue()); + + switch (r) { + case CELL_INSIDE_QUERY: + collector.countNode((int) pointTree.size()); + debug.visitInner(); + break; + case CELL_CROSSES_QUERY: + if (pointTree.moveToChild()) { + do { + intersectWithRanges(visitor, pointTree, collector, debug); + } while (pointTree.moveToSibling()); + pointTree.moveToParent(); + } else { + pointTree.visitDocValues(visitor); + debug.visitLeaf(); + } + break; + case CELL_OUTSIDE_QUERY: + } + } + + private static PointValues.IntersectVisitor getIntersectVisitor(RangeCollectorForPointTree collector) { + return new PointValues.IntersectVisitor() { + @Override + public void visit(int docID) { + // this branch should be unreachable + throw new UnsupportedOperationException( + "This IntersectVisitor does not perform any actions on a " + "docID=" + docID + " node being visited" + ); + } + + @Override + public void visit(int docID, byte[] packedValue) throws IOException { + visitPoints(packedValue, collector::count); + } + + @Override + public void visit(DocIdSetIterator iterator, byte[] packedValue) throws IOException { + visitPoints(packedValue, () -> { + for (int doc = iterator.nextDoc(); doc != NO_MORE_DOCS; doc = iterator.nextDoc()) { + collector.count(); + } + }); + } + + private void visitPoints(byte[] packedValue, CheckedRunnable collect) throws IOException { + if (!collector.withinUpperBound(packedValue)) { + collector.finalizePreviousRange(); + if (collector.iterateRangeEnd(packedValue)) { + throw new CollectionTerminatedException(); + } + } + + if (collector.withinRange(packedValue)) { + collect.run(); + } + } + + @Override + public PointValues.Relation compare(byte[] minPackedValue, byte[] maxPackedValue) { + // try to find the first range that may collect values from this cell + if (!collector.withinUpperBound(minPackedValue)) { + collector.finalizePreviousRange(); + if (collector.iterateRangeEnd(minPackedValue)) { + throw new CollectionTerminatedException(); + } + } + // after the loop, min < upper + // cell could be outside [min max] lower + if (!collector.withinLowerBound(maxPackedValue)) { + return PointValues.Relation.CELL_OUTSIDE_QUERY; + } + if (collector.withinRange(minPackedValue) && collector.withinRange(maxPackedValue)) { + return PointValues.Relation.CELL_INSIDE_QUERY; + } + return PointValues.Relation.CELL_CROSSES_QUERY; + } + }; + } + + private static class RangeCollectorForPointTree { + private final BiConsumer incrementRangeDocCount; + private int counter = 0; + + private final Ranges ranges; + private int activeIndex; + + private int visitedRange = 0; + private final int maxNumNonZeroRange; + + public RangeCollectorForPointTree( + BiConsumer incrementRangeDocCount, + int maxNumNonZeroRange, + Ranges ranges, + int activeIndex + ) { + this.incrementRangeDocCount = incrementRangeDocCount; + this.maxNumNonZeroRange = maxNumNonZeroRange; + this.ranges = ranges; + this.activeIndex = activeIndex; + } + + private void count() { + counter++; + } + + private void countNode(int count) { + counter += count; + } + + private void finalizePreviousRange() { + if (counter > 0) { + incrementRangeDocCount.accept(activeIndex, counter); + counter = 0; + } + } + + /** + * @return true when iterator exhausted or collect enough non-zero ranges + */ + private boolean iterateRangeEnd(byte[] value) { + // the new value may not be contiguous to the previous one + // so try to find the first next range that cross the new value + while (!withinUpperBound(value)) { + if (++activeIndex >= ranges.size) { + return true; + } + } + visitedRange++; + return visitedRange > maxNumNonZeroRange; + } + + private boolean withinLowerBound(byte[] value) { + return Ranges.withinLowerBound(value, ranges.lowers[activeIndex]); + } + + private boolean withinUpperBound(byte[] value) { + return Ranges.withinUpperBound(value, ranges.uppers[activeIndex]); + } + + private boolean withinRange(byte[] value) { + return withinLowerBound(value) && withinUpperBound(value); + } + } +} diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/RangeAggregatorBridge.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/RangeAggregatorBridge.java new file mode 100644 index 0000000000000..b590a444c8b04 --- /dev/null +++ b/server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/RangeAggregatorBridge.java @@ -0,0 +1,96 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.search.aggregations.bucket.filterrewrite; + +import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.index.PointValues; +import org.opensearch.index.mapper.MappedFieldType; +import org.opensearch.index.mapper.NumericPointEncoder; +import org.opensearch.search.aggregations.bucket.range.RangeAggregator; +import org.opensearch.search.aggregations.support.ValuesSource; +import org.opensearch.search.aggregations.support.ValuesSourceConfig; + +import java.io.IOException; +import java.util.function.BiConsumer; +import java.util.function.Function; + +import static org.opensearch.search.aggregations.bucket.filterrewrite.PointTreeTraversal.multiRangesTraverse; + +/** + * For range aggregation + */ +public abstract class RangeAggregatorBridge extends AggregatorBridge { + + protected boolean canOptimize(ValuesSourceConfig config, RangeAggregator.Range[] ranges) { + if (config.fieldType() == null) return false; + MappedFieldType fieldType = config.fieldType(); + assert fieldType != null; + if (fieldType.isSearchable() == false || !(fieldType instanceof NumericPointEncoder)) return false; + + if (config.script() == null && config.missing() == null) { + if (config.getValuesSource() instanceof ValuesSource.Numeric.FieldData) { + // ranges are already sorted by from and then to + // we want ranges not overlapping with each other + double prevTo = ranges[0].getTo(); + for (int i = 1; i < ranges.length; i++) { + if (prevTo > ranges[i].getFrom()) { + return false; + } + prevTo = ranges[i].getTo(); + } + this.fieldType = config.fieldType(); + return true; + } + } + return false; + } + + protected void buildRanges(RangeAggregator.Range[] ranges) { + assert fieldType instanceof NumericPointEncoder; + NumericPointEncoder numericPointEncoder = (NumericPointEncoder) fieldType; + byte[][] lowers = new byte[ranges.length][]; + byte[][] uppers = new byte[ranges.length][]; + for (int i = 0; i < ranges.length; i++) { + double rangeMin = ranges[i].getFrom(); + double rangeMax = ranges[i].getTo(); + byte[] lower = numericPointEncoder.encodePoint(rangeMin); + byte[] upper = numericPointEncoder.encodePoint(rangeMax); + lowers[i] = lower; + uppers[i] = upper; + } + + setRanges.accept(new Ranges(lowers, uppers)); + } + + @Override + final Ranges tryBuildRangesFromSegment(LeafReaderContext leaf) { + throw new UnsupportedOperationException("Range aggregation should not build ranges at segment level"); + } + + @Override + final FilterRewriteOptimizationContext.DebugInfo tryOptimize( + PointValues values, + BiConsumer incrementDocCount, + Ranges ranges + ) throws IOException { + int size = Integer.MAX_VALUE; + + BiConsumer incrementFunc = (activeIndex, docCount) -> { + long bucketOrd = bucketOrdProducer().apply(activeIndex); + incrementDocCount.accept(bucketOrd, (long) docCount); + }; + + return multiRangesTraverse(values.getPointTree(), ranges, incrementFunc, size); + } + + /** + * Provides a function to produce bucket ordinals from index of the corresponding range in the range array + */ + protected abstract Function bucketOrdProducer(); +} diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/Ranges.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/Ranges.java new file mode 100644 index 0000000000000..2819778ce215b --- /dev/null +++ b/server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/Ranges.java @@ -0,0 +1,57 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.search.aggregations.bucket.filterrewrite; + +import org.apache.lucene.util.ArrayUtil; + +/** + * Internal ranges representation for the filter rewrite optimization + */ +final class Ranges { + byte[][] lowers; // inclusive + byte[][] uppers; // exclusive + int size; + int byteLen; + static ArrayUtil.ByteArrayComparator comparator; + + Ranges(byte[][] lowers, byte[][] uppers) { + this.lowers = lowers; + this.uppers = uppers; + assert lowers.length == uppers.length; + this.size = lowers.length; + this.byteLen = lowers[0].length; + comparator = ArrayUtil.getUnsignedComparator(byteLen); + } + + public int firstRangeIndex(byte[] globalMin, byte[] globalMax) { + if (compareByteValue(lowers[0], globalMax) > 0) { + return -1; + } + int i = 0; + while (compareByteValue(uppers[i], globalMin) <= 0) { + i++; + if (i >= size) { + return -1; + } + } + return i; + } + + public static int compareByteValue(byte[] value1, byte[] value2) { + return comparator.compare(value1, 0, value2, 0); + } + + public static boolean withinLowerBound(byte[] value, byte[] lowerBound) { + return compareByteValue(value, lowerBound) >= 0; + } + + public static boolean withinUpperBound(byte[] value, byte[] upperBound) { + return compareByteValue(value, upperBound) < 0; + } +} diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/package-info.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/package-info.java new file mode 100644 index 0000000000000..bbd0a8db6cbb6 --- /dev/null +++ b/server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/package-info.java @@ -0,0 +1,19 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +/** + * This package contains filter rewrite optimization for range-type aggregations + *

+ * The idea is to + *

    + *
  • figure out the "ranges" from the aggregation
    • + *
  • leverage the ranges and bkd index to get the result of each range bucket quickly
    • + *
+ * More details in https://github.com/opensearch-project/OpenSearch/pull/14464 + */ +package org.opensearch.search.aggregations.bucket.filterrewrite; diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/histogram/AutoDateHistogramAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/histogram/AutoDateHistogramAggregator.java index d13d575a9d696..f3a36b4882d19 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/bucket/histogram/AutoDateHistogramAggregator.java +++ b/server/src/main/java/org/opensearch/search/aggregations/bucket/histogram/AutoDateHistogramAggregator.java @@ -42,7 +42,6 @@ import org.opensearch.common.util.IntArray; import org.opensearch.common.util.LongArray; import org.opensearch.core.common.util.ByteArray; -import org.opensearch.index.mapper.MappedFieldType; import org.opensearch.search.DocValueFormat; import org.opensearch.search.aggregations.Aggregator; import org.opensearch.search.aggregations.AggregatorFactories; @@ -53,8 +52,9 @@ import org.opensearch.search.aggregations.LeafBucketCollectorBase; import org.opensearch.search.aggregations.bucket.DeferableBucketAggregator; import org.opensearch.search.aggregations.bucket.DeferringBucketCollector; -import org.opensearch.search.aggregations.bucket.FastFilterRewriteHelper; import org.opensearch.search.aggregations.bucket.MergingBucketsDeferringCollector; +import org.opensearch.search.aggregations.bucket.filterrewrite.DateHistogramAggregatorBridge; +import org.opensearch.search.aggregations.bucket.filterrewrite.FilterRewriteOptimizationContext; import org.opensearch.search.aggregations.bucket.histogram.AutoDateHistogramAggregationBuilder.RoundingInfo; import org.opensearch.search.aggregations.bucket.terms.LongKeyedBucketOrds; import org.opensearch.search.aggregations.support.ValuesSource; @@ -64,11 +64,12 @@ import java.io.IOException; import java.util.Collections; import java.util.Map; -import java.util.Objects; import java.util.function.BiConsumer; import java.util.function.Function; import java.util.function.LongToIntFunction; +import static org.opensearch.search.aggregations.bucket.filterrewrite.DateHistogramAggregatorBridge.segmentMatchAll; + /** * An aggregator for date values that attempts to return a specific number of * buckets, reconfiguring how it rounds dates to buckets on the fly as new @@ -135,7 +136,7 @@ static AutoDateHistogramAggregator build( protected int roundingIdx; protected Rounding.Prepared preparedRounding; - private final FastFilterRewriteHelper.FastFilterContext fastFilterContext; + private final FilterRewriteOptimizationContext filterRewriteOptimizationContext; private AutoDateHistogramAggregator( String name, @@ -158,53 +159,52 @@ private AutoDateHistogramAggregator( this.roundingPreparer = roundingPreparer; this.preparedRounding = prepareRounding(0); - fastFilterContext = new FastFilterRewriteHelper.FastFilterContext( - context, - new AutoHistogramAggregationType( - valuesSourceConfig.fieldType(), - valuesSourceConfig.missing() != null, - valuesSourceConfig.script() != null - ) - ); - if (fastFilterContext.isRewriteable(parent, subAggregators.length)) { - fastFilterContext.buildRanges(Objects.requireNonNull(valuesSourceConfig.fieldType())); - } - } - - private class AutoHistogramAggregationType extends FastFilterRewriteHelper.AbstractDateHistogramAggregationType { + DateHistogramAggregatorBridge bridge = new DateHistogramAggregatorBridge() { + @Override + protected boolean canOptimize() { + return canOptimize(valuesSourceConfig); + } - public AutoHistogramAggregationType(MappedFieldType fieldType, boolean missing, boolean hasScript) { - super(fieldType, missing, hasScript); - } + @Override + protected void prepare() throws IOException { + buildRanges(context); + } - @Override - protected Rounding getRounding(final long low, final long high) { - // max - min / targetBuckets = bestDuration - // find the right innerInterval this bestDuration belongs to - // since we cannot exceed targetBuckets, bestDuration should go up, - // so the right innerInterval should be an upper bound - long bestDuration = (high - low) / targetBuckets; - // reset so this function is idempotent - roundingIdx = 0; - while (roundingIdx < roundingInfos.length - 1) { - final RoundingInfo curRoundingInfo = roundingInfos[roundingIdx]; - final int temp = curRoundingInfo.innerIntervals[curRoundingInfo.innerIntervals.length - 1]; - // If the interval duration is covered by the maximum inner interval, - // we can start with this outer interval for creating the buckets - if (bestDuration <= temp * curRoundingInfo.roughEstimateDurationMillis) { - break; + @Override + protected Rounding getRounding(final long low, final long high) { + // max - min / targetBuckets = bestDuration + // find the right innerInterval this bestDuration belongs to + // since we cannot exceed targetBuckets, bestDuration should go up, + // so the right innerInterval should be an upper bound + long bestDuration = (high - low) / targetBuckets; + // reset so this function is idempotent + roundingIdx = 0; + while (roundingIdx < roundingInfos.length - 1) { + final RoundingInfo curRoundingInfo = roundingInfos[roundingIdx]; + final int temp = curRoundingInfo.innerIntervals[curRoundingInfo.innerIntervals.length - 1]; + // If the interval duration is covered by the maximum inner interval, + // we can start with this outer interval for creating the buckets + if (bestDuration <= temp * curRoundingInfo.roughEstimateDurationMillis) { + break; + } + roundingIdx++; } - roundingIdx++; + + preparedRounding = prepareRounding(roundingIdx); + return roundingInfos[roundingIdx].rounding; } - preparedRounding = prepareRounding(roundingIdx); - return roundingInfos[roundingIdx].rounding; - } + @Override + protected Prepared getRoundingPrepared() { + return preparedRounding; + } - @Override - protected Prepared getRoundingPrepared() { - return preparedRounding; - } + @Override + protected Function bucketOrdProducer() { + return (key) -> getBucketOrds().add(0, preparedRounding.round((long) key)); + } + }; + filterRewriteOptimizationContext = new FilterRewriteOptimizationContext(bridge, parent, subAggregators.length, context); } protected abstract LongKeyedBucketOrds getBucketOrds(); @@ -236,11 +236,7 @@ public final LeafBucketCollector getLeafCollector(LeafReaderContext ctx, LeafBuc return LeafBucketCollector.NO_OP_COLLECTOR; } - boolean optimized = fastFilterContext.tryFastFilterAggregation( - ctx, - this::incrementBucketDocCount, - (key) -> getBucketOrds().add(0, preparedRounding.round((long) key)) - ); + boolean optimized = filterRewriteOptimizationContext.tryOptimize(ctx, this::incrementBucketDocCount, segmentMatchAll(context, ctx)); if (optimized) throw new CollectionTerminatedException(); final SortedNumericDocValues values = valuesSource.longValues(ctx); @@ -308,12 +304,7 @@ protected final void merge(long[] mergeMap, long newNumBuckets) { @Override public void collectDebugInfo(BiConsumer add) { super.collectDebugInfo(add); - if (fastFilterContext.optimizedSegments > 0) { - add.accept("optimized_segments", fastFilterContext.optimizedSegments); - add.accept("unoptimized_segments", fastFilterContext.segments - fastFilterContext.optimizedSegments); - add.accept("leaf_visited", fastFilterContext.leaf); - add.accept("inner_visited", fastFilterContext.inner); - } + filterRewriteOptimizationContext.populateDebugInfo(add); } /** diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/histogram/DateHistogramAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/histogram/DateHistogramAggregator.java index 4b84797c18922..96a49bc3fd5f6 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/bucket/histogram/DateHistogramAggregator.java +++ b/server/src/main/java/org/opensearch/search/aggregations/bucket/histogram/DateHistogramAggregator.java @@ -39,7 +39,6 @@ import org.opensearch.common.Nullable; import org.opensearch.common.Rounding; import org.opensearch.common.lease.Releasables; -import org.opensearch.index.mapper.MappedFieldType; import org.opensearch.search.DocValueFormat; import org.opensearch.search.aggregations.Aggregator; import org.opensearch.search.aggregations.AggregatorFactories; @@ -49,7 +48,8 @@ import org.opensearch.search.aggregations.LeafBucketCollector; import org.opensearch.search.aggregations.LeafBucketCollectorBase; import org.opensearch.search.aggregations.bucket.BucketsAggregator; -import org.opensearch.search.aggregations.bucket.FastFilterRewriteHelper; +import org.opensearch.search.aggregations.bucket.filterrewrite.DateHistogramAggregatorBridge; +import org.opensearch.search.aggregations.bucket.filterrewrite.FilterRewriteOptimizationContext; import org.opensearch.search.aggregations.bucket.terms.LongKeyedBucketOrds; import org.opensearch.search.aggregations.support.ValuesSource; import org.opensearch.search.aggregations.support.ValuesSourceConfig; @@ -58,8 +58,10 @@ import java.io.IOException; import java.util.Collections; import java.util.Map; -import java.util.Objects; import java.util.function.BiConsumer; +import java.util.function.Function; + +import static org.opensearch.search.aggregations.bucket.filterrewrite.DateHistogramAggregatorBridge.segmentMatchAll; /** * An aggregator for date values. Every date is rounded down using a configured @@ -84,7 +86,7 @@ class DateHistogramAggregator extends BucketsAggregator implements SizedBucketAg private final LongBounds hardBounds; private final LongKeyedBucketOrds bucketOrds; - private final FastFilterRewriteHelper.FastFilterContext fastFilterContext; + private final FilterRewriteOptimizationContext filterRewriteOptimizationContext; DateHistogramAggregator( String name, @@ -117,35 +119,38 @@ class DateHistogramAggregator extends BucketsAggregator implements SizedBucketAg bucketOrds = LongKeyedBucketOrds.build(context.bigArrays(), cardinality); - fastFilterContext = new FastFilterRewriteHelper.FastFilterContext( - context, - new DateHistogramAggregationType( - valuesSourceConfig.fieldType(), - valuesSourceConfig.missing() != null, - valuesSourceConfig.script() != null, - hardBounds - ) - ); - if (fastFilterContext.isRewriteable(parent, subAggregators.length)) { - fastFilterContext.buildRanges(Objects.requireNonNull(valuesSourceConfig.fieldType())); - } - } + DateHistogramAggregatorBridge bridge = new DateHistogramAggregatorBridge() { + @Override + protected boolean canOptimize() { + return canOptimize(valuesSourceConfig); + } - private class DateHistogramAggregationType extends FastFilterRewriteHelper.AbstractDateHistogramAggregationType { + @Override + protected void prepare() throws IOException { + buildRanges(context); + } - public DateHistogramAggregationType(MappedFieldType fieldType, boolean missing, boolean hasScript, LongBounds hardBounds) { - super(fieldType, missing, hasScript, hardBounds); - } + @Override + protected Rounding getRounding(long low, long high) { + return rounding; + } - @Override - protected Rounding getRounding(long low, long high) { - return rounding; - } + @Override + protected Rounding.Prepared getRoundingPrepared() { + return preparedRounding; + } - @Override - protected Rounding.Prepared getRoundingPrepared() { - return preparedRounding; - } + @Override + protected long[] processHardBounds(long[] bounds) { + return super.processHardBounds(bounds, hardBounds); + } + + @Override + protected Function bucketOrdProducer() { + return (key) -> bucketOrds.add(0, preparedRounding.round((long) key)); + } + }; + filterRewriteOptimizationContext = new FilterRewriteOptimizationContext(bridge, parent, subAggregators.length, context); } @Override @@ -162,11 +167,7 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, LeafBucketCol return LeafBucketCollector.NO_OP_COLLECTOR; } - boolean optimized = fastFilterContext.tryFastFilterAggregation( - ctx, - this::incrementBucketDocCount, - (key) -> bucketOrds.add(0, preparedRounding.round((long) key)) - ); + boolean optimized = filterRewriteOptimizationContext.tryOptimize(ctx, this::incrementBucketDocCount, segmentMatchAll(context, ctx)); if (optimized) throw new CollectionTerminatedException(); SortedNumericDocValues values = valuesSource.longValues(ctx); @@ -253,12 +254,7 @@ public void doClose() { @Override public void collectDebugInfo(BiConsumer add) { add.accept("total_buckets", bucketOrds.size()); - if (fastFilterContext.optimizedSegments > 0) { - add.accept("optimized_segments", fastFilterContext.optimizedSegments); - add.accept("unoptimized_segments", fastFilterContext.segments - fastFilterContext.optimizedSegments); - add.accept("leaf_visited", fastFilterContext.leaf); - add.accept("inner_visited", fastFilterContext.inner); - } + filterRewriteOptimizationContext.populateDebugInfo(add); } /** diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/range/RangeAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/range/RangeAggregator.java index 2ba2b06514de1..17461f228e993 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/bucket/range/RangeAggregator.java +++ b/server/src/main/java/org/opensearch/search/aggregations/bucket/range/RangeAggregator.java @@ -55,7 +55,8 @@ import org.opensearch.search.aggregations.LeafBucketCollectorBase; import org.opensearch.search.aggregations.NonCollectingAggregator; import org.opensearch.search.aggregations.bucket.BucketsAggregator; -import org.opensearch.search.aggregations.bucket.FastFilterRewriteHelper; +import org.opensearch.search.aggregations.bucket.filterrewrite.FilterRewriteOptimizationContext; +import org.opensearch.search.aggregations.bucket.filterrewrite.RangeAggregatorBridge; import org.opensearch.search.aggregations.support.ValuesSource; import org.opensearch.search.aggregations.support.ValuesSourceConfig; import org.opensearch.search.internal.SearchContext; @@ -66,6 +67,7 @@ import java.util.Map; import java.util.Objects; import java.util.function.BiConsumer; +import java.util.function.Function; import static org.opensearch.core.xcontent.ConstructingObjectParser.optionalConstructorArg; @@ -249,7 +251,7 @@ public boolean equals(Object obj) { final double[] maxTo; - private final FastFilterRewriteHelper.FastFilterContext fastFilterContext; + private final FilterRewriteOptimizationContext filterRewriteOptimizationContext; public RangeAggregator( String name, @@ -279,13 +281,23 @@ public RangeAggregator( maxTo[i] = Math.max(this.ranges[i].to, maxTo[i - 1]); } - fastFilterContext = new FastFilterRewriteHelper.FastFilterContext( - context, - new FastFilterRewriteHelper.RangeAggregationType(config, ranges) - ); - if (fastFilterContext.isRewriteable(parent, subAggregators.length)) { - fastFilterContext.buildRanges(Objects.requireNonNull(config.fieldType())); - } + RangeAggregatorBridge bridge = new RangeAggregatorBridge() { + @Override + protected boolean canOptimize() { + return canOptimize(config, ranges); + } + + @Override + protected void prepare() { + buildRanges(ranges); + } + + @Override + protected Function bucketOrdProducer() { + return (activeIndex) -> subBucketOrdinal(0, (int) activeIndex); + } + }; + filterRewriteOptimizationContext = new FilterRewriteOptimizationContext(bridge, parent, subAggregators.length, context); } @Override @@ -298,11 +310,7 @@ public ScoreMode scoreMode() { @Override public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBucketCollector sub) throws IOException { - boolean optimized = fastFilterContext.tryFastFilterAggregation( - ctx, - this::incrementBucketDocCount, - (activeIndex) -> subBucketOrdinal(0, (int) activeIndex) - ); + boolean optimized = filterRewriteOptimizationContext.tryOptimize(ctx, this::incrementBucketDocCount, false); if (optimized) throw new CollectionTerminatedException(); final SortedNumericDoubleValues values = valuesSource.doubleValues(ctx); @@ -452,11 +460,6 @@ public InternalAggregation buildEmptyAggregation() { @Override public void collectDebugInfo(BiConsumer add) { super.collectDebugInfo(add); - if (fastFilterContext.optimizedSegments > 0) { - add.accept("optimized_segments", fastFilterContext.optimizedSegments); - add.accept("unoptimized_segments", fastFilterContext.segments - fastFilterContext.optimizedSegments); - add.accept("leaf_visited", fastFilterContext.leaf); - add.accept("inner_visited", fastFilterContext.inner); - } + filterRewriteOptimizationContext.populateDebugInfo(add); } } From 8f95735f0b75eb3fca117f82ebeda5fd32c55f6b Mon Sep 17 00:00:00 2001 From: gaobinlong Date: Fri, 9 Aug 2024 20:01:35 +0800 Subject: [PATCH 11/40] Fix flaky test in {yaml=ingest/70_bulk} (#15172) Signed-off-by: Gao Binlong --- .../resources/rest-api-spec/test/ingest/70_bulk.yml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/modules/ingest-common/src/yamlRestTest/resources/rest-api-spec/test/ingest/70_bulk.yml b/modules/ingest-common/src/yamlRestTest/resources/rest-api-spec/test/ingest/70_bulk.yml index 47cc80d6df310..ecd56ea7f277e 100644 --- a/modules/ingest-common/src/yamlRestTest/resources/rest-api-spec/test/ingest/70_bulk.yml +++ b/modules/ingest-common/src/yamlRestTest/resources/rest-api-spec/test/ingest/70_bulk.yml @@ -178,7 +178,10 @@ teardown: - skip: version: " - 2.15.99" reason: "fixed in 2.16.0" + features: allowed_warnings - do: + allowed_warnings: + - "index template [test_for_bulk_upsert_index_template] has index patterns [test_bulk_upsert_*] matching patterns from existing older templates [global] with patterns (global => [*]); this template [test_for_bulk_upsert_index_template] will take precedence during new index creation" indices.put_index_template: name: test_for_bulk_upsert_index_template body: From 875b603ccead43b5302195c71208e758af1848c2 Mon Sep 17 00:00:00 2001 From: Gaurav Bafna <85113518+gbbafna@users.noreply.github.com> Date: Fri, 9 Aug 2024 17:35:48 +0530 Subject: [PATCH 12/40] =?UTF-8?q?Flaky=20test=20:=20Don't=20use=20async=20?= =?UTF-8?q?repo=20for=20SplitIndex=20and=20wait=20for=20translo=E2=80=A6?= =?UTF-8?q?=20(#15107)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Flaky test : Don't use async repo for SplitIndex and wait for translog file deletion Signed-off-by: Gaurav Bafna --- .../indices/create/RemoteSplitIndexIT.java | 24 +++++++++++++++++++ .../RemoteStoreBaseIntegTestCase.java | 2 +- .../remotestore/RemoteStoreStatsIT.java | 18 ++++++++------ 3 files changed, 36 insertions(+), 8 deletions(-) diff --git a/server/src/internalClusterTest/java/org/opensearch/action/admin/indices/create/RemoteSplitIndexIT.java b/server/src/internalClusterTest/java/org/opensearch/action/admin/indices/create/RemoteSplitIndexIT.java index dc3c8793a93f6..928c9e33e19cb 100644 --- a/server/src/internalClusterTest/java/org/opensearch/action/admin/indices/create/RemoteSplitIndexIT.java +++ b/server/src/internalClusterTest/java/org/opensearch/action/admin/indices/create/RemoteSplitIndexIT.java @@ -46,6 +46,7 @@ import org.opensearch.action.index.IndexRequest; import org.opensearch.action.index.IndexRequestBuilder; import org.opensearch.action.search.SearchResponse; +import org.opensearch.action.support.IndicesOptions; import org.opensearch.client.Client; import org.opensearch.cluster.ClusterState; import org.opensearch.cluster.metadata.IndexMetadata; @@ -69,12 +70,15 @@ import org.opensearch.remotestore.RemoteStoreBaseIntegTestCase; import org.opensearch.test.OpenSearchIntegTestCase; import org.opensearch.test.VersionUtils; +import org.junit.After; +import org.junit.Before; import java.io.IOException; import java.io.UncheckedIOException; import java.util.Arrays; import java.util.HashSet; import java.util.Set; +import java.util.concurrent.TimeUnit; import java.util.function.BiFunction; import java.util.stream.IntStream; @@ -89,12 +93,32 @@ @OpenSearchIntegTestCase.ClusterScope(scope = OpenSearchIntegTestCase.Scope.TEST, numDataNodes = 0) public class RemoteSplitIndexIT extends RemoteStoreBaseIntegTestCase { + @Before + public void setup() { + asyncUploadMockFsRepo = false; + } @Override protected boolean forbidPrivateIndexSettings() { return false; } + @After + public void cleanUp() throws Exception { + // Delete is async. + assertAcked( + client().admin().indices().prepareDelete("*").setIndicesOptions(IndicesOptions.LENIENT_EXPAND_OPEN_CLOSED_HIDDEN).get() + ); + assertBusy(() -> { + try { + assertEquals(0, getFileCount(translogRepoPath)); + } catch (IOException e) { + fail(); + } + }, 30, TimeUnit.SECONDS); + super.teardown(); + } + public Settings indexSettings() { return Settings.builder() .put(super.indexSettings()) diff --git a/server/src/internalClusterTest/java/org/opensearch/remotestore/RemoteStoreBaseIntegTestCase.java b/server/src/internalClusterTest/java/org/opensearch/remotestore/RemoteStoreBaseIntegTestCase.java index 63a9451a27a12..f83ae3e0ca820 100644 --- a/server/src/internalClusterTest/java/org/opensearch/remotestore/RemoteStoreBaseIntegTestCase.java +++ b/server/src/internalClusterTest/java/org/opensearch/remotestore/RemoteStoreBaseIntegTestCase.java @@ -313,7 +313,7 @@ public void assertRemoteStoreRepositoryOnAllNodes(String repositoryName) { } } - public static int getFileCount(Path path) throws Exception { + public static int getFileCount(Path path) throws IOException { final AtomicInteger filesExisting = new AtomicInteger(0); Files.walkFileTree(path, new SimpleFileVisitor<>() { @Override diff --git a/server/src/internalClusterTest/java/org/opensearch/remotestore/RemoteStoreStatsIT.java b/server/src/internalClusterTest/java/org/opensearch/remotestore/RemoteStoreStatsIT.java index 31c73e2fc03ae..86d586cd17146 100644 --- a/server/src/internalClusterTest/java/org/opensearch/remotestore/RemoteStoreStatsIT.java +++ b/server/src/internalClusterTest/java/org/opensearch/remotestore/RemoteStoreStatsIT.java @@ -250,7 +250,7 @@ public void testStatsResponseFromLocalNode() { } } - @TestLogging(reason = "Getting trace logs from remote store package", value = "org.opensearch.remotestore:TRACE") + @TestLogging(reason = "Getting trace logs from remote store package", value = "org.opensearch.index.shard:TRACE") public void testDownloadStatsCorrectnessSinglePrimarySingleReplica() throws Exception { setup(); // Scenario: @@ -280,11 +280,13 @@ public void testDownloadStatsCorrectnessSinglePrimarySingleReplica() throws Exce .get(0) .getSegmentStats(); logger.info( - "Zero state primary stats: {}ms refresh time lag, {}b bytes lag, {}b upload bytes started and {}b upload bytes failed.", + "Zero state primary stats: {}ms refresh time lag, {}b bytes lag, {}b upload bytes started, {}b upload bytes failed , {} uploads succeeded, {} upload byes succeeded.", zeroStatePrimaryStats.refreshTimeLagMs, zeroStatePrimaryStats.bytesLag, zeroStatePrimaryStats.uploadBytesStarted, - zeroStatePrimaryStats.uploadBytesFailed + zeroStatePrimaryStats.uploadBytesFailed, + zeroStatePrimaryStats.totalUploadsSucceeded, + zeroStatePrimaryStats.uploadBytesSucceeded ); assertTrue( zeroStatePrimaryStats.totalUploadsStarted == zeroStatePrimaryStats.totalUploadsSucceeded @@ -348,7 +350,7 @@ public void testDownloadStatsCorrectnessSinglePrimarySingleReplica() throws Exce } } - @TestLogging(reason = "Getting trace logs from remote store package", value = "org.opensearch.remotestore:TRACE") + @TestLogging(reason = "Getting trace logs from remote store package", value = "org.opensearch.index.shard:TRACE") public void testDownloadStatsCorrectnessSinglePrimaryMultipleReplicaShards() throws Exception { setup(); // Scenario: @@ -382,11 +384,13 @@ public void testDownloadStatsCorrectnessSinglePrimaryMultipleReplicaShards() thr .get(0) .getSegmentStats(); logger.info( - "Zero state primary stats: {}ms refresh time lag, {}b bytes lag, {}b upload bytes started and {}b upload bytes failed.", + "Zero state primary stats: {}ms refresh time lag, {}b bytes lag, {}b upload bytes started, {}b upload bytes failed , {} uploads succeeded, {} upload byes succeeded.", zeroStatePrimaryStats.refreshTimeLagMs, zeroStatePrimaryStats.bytesLag, zeroStatePrimaryStats.uploadBytesStarted, - zeroStatePrimaryStats.uploadBytesFailed + zeroStatePrimaryStats.uploadBytesFailed, + zeroStatePrimaryStats.totalUploadsSucceeded, + zeroStatePrimaryStats.uploadBytesSucceeded ); assertTrue( zeroStatePrimaryStats.totalUploadsStarted == zeroStatePrimaryStats.totalUploadsSucceeded @@ -617,7 +621,7 @@ public void testNonZeroPrimaryStatsOnNewlyCreatedIndexWithZeroDocs() throws Exce } assertZeroTranslogDownloadStats(translogStats); }); - }, 5, TimeUnit.SECONDS); + }, 10, TimeUnit.SECONDS); } public void testStatsCorrectnessOnFailover() { From 3b358e3ad6feca8e1e7f6a39f46c33a4db63e9dc Mon Sep 17 00:00:00 2001 From: Andriy Redko Date: Fri, 9 Aug 2024 13:55:01 -0400 Subject: [PATCH 13/40] Fix AzureBlobContainerRetriesTests flaky tests (#15184) --- .../repositories/azure/AzureBlobContainerRetriesTests.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/plugins/repository-azure/src/test/java/org/opensearch/repositories/azure/AzureBlobContainerRetriesTests.java b/plugins/repository-azure/src/test/java/org/opensearch/repositories/azure/AzureBlobContainerRetriesTests.java index 71ffd0fd959f1..acb3547b23f0d 100644 --- a/plugins/repository-azure/src/test/java/org/opensearch/repositories/azure/AzureBlobContainerRetriesTests.java +++ b/plugins/repository-azure/src/test/java/org/opensearch/repositories/azure/AzureBlobContainerRetriesTests.java @@ -157,7 +157,7 @@ private BlobContainer createBlobContainer(final int maxRetries) { + "/"; clientSettings.put(ENDPOINT_SUFFIX_SETTING.getConcreteSettingForNamespace(clientName).getKey(), endpoint); clientSettings.put(MAX_RETRIES_SETTING.getConcreteSettingForNamespace(clientName).getKey(), maxRetries); - clientSettings.put(TIMEOUT_SETTING.getConcreteSettingForNamespace(clientName).getKey(), TimeValue.timeValueMillis(2000)); + clientSettings.put(TIMEOUT_SETTING.getConcreteSettingForNamespace(clientName).getKey(), TimeValue.timeValueMillis(5000)); final MockSecureSettings secureSettings = new MockSecureSettings(); secureSettings.setString(ACCOUNT_SETTING.getConcreteSettingForNamespace(clientName).getKey(), "account"); From 252742b4c97ce1bc77a3589cec30ed6a66787b8c Mon Sep 17 00:00:00 2001 From: gaobinlong Date: Sat, 10 Aug 2024 03:02:38 +0800 Subject: [PATCH 14/40] Fix array_index_out_of_bounds_exception when indexing documents with field name containing only dot (#15126) * Fix array_index_out_of_bounds_exception when indexing documents with field name containing only dot Signed-off-by: Gao Binlong * Modify change log Signed-off-by: Gao Binlong * Optimize error message Signed-off-by: Gao Binlong --------- Signed-off-by: Gao Binlong --- CHANGELOG.md | 1 + .../test/index/120_field_name.yml | 38 +++++++++++++++++++ .../index/mapper/DocumentParser.java | 3 ++ .../index/mapper/DocumentParserTests.java | 32 ++++++++++++++++ 4 files changed, 74 insertions(+) create mode 100644 rest-api-spec/src/main/resources/rest-api-spec/test/index/120_field_name.yml diff --git a/CHANGELOG.md b/CHANGELOG.md index 1964e456acda0..7fd56693e1ea6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -42,6 +42,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), - Fix NPE when bulk ingest with empty pipeline ([#15033](https://github.com/opensearch-project/OpenSearch/pull/15033)) - Fix missing value of FieldSort for unsigned_long ([#14963](https://github.com/opensearch-project/OpenSearch/pull/14963)) - Fix delete index template failed when the index template matches a data stream but is unused ([#15080](https://github.com/opensearch-project/OpenSearch/pull/15080)) +- Fix array_index_out_of_bounds_exception when indexing documents with field name containing only dot ([#15126](https://github.com/opensearch-project/OpenSearch/pull/15126)) ### Security diff --git a/rest-api-spec/src/main/resources/rest-api-spec/test/index/120_field_name.yml b/rest-api-spec/src/main/resources/rest-api-spec/test/index/120_field_name.yml new file mode 100644 index 0000000000000..dae1da9963b0d --- /dev/null +++ b/rest-api-spec/src/main/resources/rest-api-spec/test/index/120_field_name.yml @@ -0,0 +1,38 @@ +--- +"Index documents with field name containing only dot fail with an IllegalArgumentException": + - skip: + version: " - 2.99.99" + reason: "introduced in 3.0.0" + + - do: + indices.create: + index: test_1 + + - do: + catch: /field name cannot contain only the character \[.\]/ + index: + index: test_1 + id: 1 + body: { + .: bar + } + + - do: + catch: /field name cannot contain only the character \[.\]/ + index: + index: test_1 + id: 1 + body: { + ..: bar + } + + - do: + catch: /field name cannot contain only the character \[.\]/ + index: + index: test_1 + id: 1 + body: { + foo: { + .: bar + } + } diff --git a/server/src/main/java/org/opensearch/index/mapper/DocumentParser.java b/server/src/main/java/org/opensearch/index/mapper/DocumentParser.java index c6815ebe8d91a..b03026d560dbf 100644 --- a/server/src/main/java/org/opensearch/index/mapper/DocumentParser.java +++ b/server/src/main/java/org/opensearch/index/mapper/DocumentParser.java @@ -206,6 +206,9 @@ private static MapperParsingException wrapInMapperParsingException(SourceToParse private static String[] splitAndValidatePath(String fullFieldPath) { if (fullFieldPath.contains(".")) { String[] parts = fullFieldPath.split("\\."); + if (parts.length == 0) { + throw new IllegalArgumentException("field name cannot contain only the character [.]"); + } for (String part : parts) { if (Strings.hasText(part) == false) { // check if the field name contains only whitespace diff --git a/server/src/test/java/org/opensearch/index/mapper/DocumentParserTests.java b/server/src/test/java/org/opensearch/index/mapper/DocumentParserTests.java index 15e2b6649b0be..c9763d634979b 100644 --- a/server/src/test/java/org/opensearch/index/mapper/DocumentParserTests.java +++ b/server/src/test/java/org/opensearch/index/mapper/DocumentParserTests.java @@ -1700,6 +1700,38 @@ public void testDynamicFieldsStartingAndEndingWithDot() throws Exception { ); } + public void testDynamicFieldsWithOnlyDot() throws Exception { + DocumentMapper mapper = createDocumentMapper(mapping(b -> {})); + + MapperParsingException e = expectThrows(MapperParsingException.class, () -> mapper.parse(source(b -> { + b.startArray("top"); + { + b.startObject(); + { + b.startObject("inner").field(".", 2).endObject(); + } + b.endObject(); + } + b.endArray(); + }))); + + assertThat(e.getCause(), notNullValue()); + assertThat(e.getCause().getMessage(), containsString("field name cannot contain only the character [.]")); + + e = expectThrows( + MapperParsingException.class, + () -> mapper.parse(source(b -> { b.startObject("..").field("foo", 2).endObject(); })) + ); + + assertThat(e.getCause(), notNullValue()); + assertThat(e.getCause().getMessage(), containsString("field name cannot contain only the character [.]")); + + e = expectThrows(MapperParsingException.class, () -> mapper.parse(source(b -> b.field(".", "1234")))); + + assertThat(e.getCause(), notNullValue()); + assertThat(e.getCause().getMessage(), containsString("field name cannot contain only the character [.]")); + } + public void testDynamicFieldsEmptyName() throws Exception { DocumentMapper mapper = createDocumentMapper(mapping(b -> {})); From 6ac92a17623ff712e8405e3e9b82f76cec50b1d0 Mon Sep 17 00:00:00 2001 From: Rishabh Singh Date: Fri, 9 Aug 2024 18:18:03 -0700 Subject: [PATCH 15/40] add required distribution key in setup-java action (#15198) Signed-off-by: Rishabh Singh --- .github/workflows/benchmark-pull-request.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/benchmark-pull-request.yml b/.github/workflows/benchmark-pull-request.yml index 0173b7e35c64d..59b3e3900823e 100644 --- a/.github/workflows/benchmark-pull-request.yml +++ b/.github/workflows/benchmark-pull-request.yml @@ -129,6 +129,7 @@ jobs: uses: actions/setup-java@v4 with: java-version: 21 + distribution: 'temurin' - name: Build and Assemble OpenSearch from PR run: | ./gradlew :distribution:archives:linux-tar:assemble -Dbuild.snapshot=false From ba7b66de37ae8b3c75f04ba84e2018d36967b1ea Mon Sep 17 00:00:00 2001 From: Andriy Redko Date: Sat, 10 Aug 2024 00:42:29 -0400 Subject: [PATCH 16/40] Fix more AzureBlobContainerRetriesTests flaky tests (#15189) --- .../repositories/azure/AzureBlobContainerRetriesTests.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/plugins/repository-azure/src/test/java/org/opensearch/repositories/azure/AzureBlobContainerRetriesTests.java b/plugins/repository-azure/src/test/java/org/opensearch/repositories/azure/AzureBlobContainerRetriesTests.java index acb3547b23f0d..970388498ee26 100644 --- a/plugins/repository-azure/src/test/java/org/opensearch/repositories/azure/AzureBlobContainerRetriesTests.java +++ b/plugins/repository-azure/src/test/java/org/opensearch/repositories/azure/AzureBlobContainerRetriesTests.java @@ -171,7 +171,7 @@ RequestRetryOptions createRetryPolicy(final AzureStorageSettings azureStorageSet return new RequestRetryOptions( RetryPolicyType.EXPONENTIAL, azureStorageSettings.getMaxRetries(), - 1, + 5, 10L, 100L, secondaryHost From 3db252541cc843eb405c6dbb7c42712d3d612146 Mon Sep 17 00:00:00 2001 From: Andriy Redko Date: Sat, 10 Aug 2024 11:14:21 -0400 Subject: [PATCH 17/40] Update 120_field_name.yml (#15192) Signed-off-by: Andriy Redko --- .../resources/rest-api-spec/test/index/120_field_name.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/rest-api-spec/src/main/resources/rest-api-spec/test/index/120_field_name.yml b/rest-api-spec/src/main/resources/rest-api-spec/test/index/120_field_name.yml index dae1da9963b0d..040e883b4a4c2 100644 --- a/rest-api-spec/src/main/resources/rest-api-spec/test/index/120_field_name.yml +++ b/rest-api-spec/src/main/resources/rest-api-spec/test/index/120_field_name.yml @@ -1,8 +1,8 @@ --- "Index documents with field name containing only dot fail with an IllegalArgumentException": - skip: - version: " - 2.99.99" - reason: "introduced in 3.0.0" + version: " - 2.16.99" + reason: "introduced in 2.17.0" - do: indices.create: From e21f6b8e768870d5426d1d1b04815a17a7a05ab3 Mon Sep 17 00:00:00 2001 From: Sarthak Aggarwal Date: Mon, 12 Aug 2024 17:51:15 +0530 Subject: [PATCH 18/40] Star Tree Min and Max Value Aggregators (#14625) --------- Signed-off-by: Sarthak Aggarwal --- .../aggregators/CountValueAggregator.java | 46 +- .../aggregators/MaxValueAggregator.java | 27 + .../aggregators/MetricAggregatorInfo.java | 13 +- .../aggregators/MinValueAggregator.java | 27 + .../StatelessDoubleValueAggregator.java | 81 +++ .../aggregators/SumValueAggregator.java | 85 ++-- .../startree/aggregators/ValueAggregator.java | 28 +- .../aggregators/ValueAggregatorFactory.java | 23 +- .../builder/AbstractDocumentsFileManager.java | 7 +- .../startree/builder/BaseStarTreeBuilder.java | 41 +- .../AbstractValueAggregatorTests.java | 77 +++ .../CountValueAggregatorTests.java | 46 +- .../aggregators/MaxValueAggregatorTests.java | 72 +++ .../aggregators/MinValueAggregatorTests.java | 71 +++ .../StaticValueAggregatorTests.java | 133 +++++ .../aggregators/SumValueAggregatorTests.java | 79 +-- .../ValueAggregatorFactoryTests.java | 27 +- .../builder/AbstractStarTreeBuilderTests.java | 475 ++++++++++++++---- 18 files changed, 1049 insertions(+), 309 deletions(-) create mode 100644 server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/MaxValueAggregator.java create mode 100644 server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/MinValueAggregator.java create mode 100644 server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/StatelessDoubleValueAggregator.java create mode 100644 server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/AbstractValueAggregatorTests.java create mode 100644 server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/MaxValueAggregatorTests.java create mode 100644 server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/MinValueAggregatorTests.java create mode 100644 server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/StaticValueAggregatorTests.java diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/CountValueAggregator.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/CountValueAggregator.java index ed159ee2efb7b..38a59d403d36b 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/CountValueAggregator.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/CountValueAggregator.java @@ -7,7 +7,6 @@ */ package org.opensearch.index.compositeindex.datacube.startree.aggregators; -import org.opensearch.index.compositeindex.datacube.MetricStat; import org.opensearch.index.compositeindex.datacube.startree.aggregators.numerictype.StarTreeNumericType; /** @@ -15,20 +14,16 @@ * * @opensearch.experimental */ -public class CountValueAggregator implements ValueAggregator { - public static final StarTreeNumericType VALUE_AGGREGATOR_TYPE = StarTreeNumericType.LONG; +class CountValueAggregator implements ValueAggregator { + public static final long DEFAULT_INITIAL_VALUE = 1L; - private StarTreeNumericType starTreeNumericType; + private final StarTreeNumericType starTreeNumericType; + private static final StarTreeNumericType VALUE_AGGREGATOR_TYPE = StarTreeNumericType.LONG; public CountValueAggregator(StarTreeNumericType starTreeNumericType) { this.starTreeNumericType = starTreeNumericType; } - @Override - public MetricStat getAggregationType() { - return MetricStat.COUNT; - } - @Override public StarTreeNumericType getAggregatedValueType() { return VALUE_AGGREGATOR_TYPE; @@ -36,34 +31,34 @@ public StarTreeNumericType getAggregatedValueType() { @Override public Long getInitialAggregatedValueForSegmentDocValue(Long segmentDocValue) { + + if (segmentDocValue == null) { + return getIdentityMetricValue(); + } + return DEFAULT_INITIAL_VALUE; } @Override public Long mergeAggregatedValueAndSegmentValue(Long value, Long segmentDocValue) { - return value + 1; + assert value != null; + if (segmentDocValue != null) { + return value + 1; + } + return value; } @Override public Long mergeAggregatedValues(Long value, Long aggregatedValue) { + if (value == null) { + value = getIdentityMetricValue(); + } + if (aggregatedValue == null) { + aggregatedValue = getIdentityMetricValue(); + } return value + aggregatedValue; } - @Override - public Long getInitialAggregatedValue(Long value) { - return value; - } - - @Override - public int getMaxAggregatedValueByteSize() { - return Long.BYTES; - } - - @Override - public Long toLongValue(Long value) { - return value; - } - @Override public Long toStarTreeNumericTypeValue(Long value) { return value; @@ -71,6 +66,7 @@ public Long toStarTreeNumericTypeValue(Long value) { @Override public Long getIdentityMetricValue() { + // in present aggregations, if the metric behind count is missing, we treat it as 0 return 0L; } } diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/MaxValueAggregator.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/MaxValueAggregator.java new file mode 100644 index 0000000000000..8e7bb44238e40 --- /dev/null +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/MaxValueAggregator.java @@ -0,0 +1,27 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ +package org.opensearch.index.compositeindex.datacube.startree.aggregators; + +import org.opensearch.index.compositeindex.datacube.startree.aggregators.numerictype.StarTreeNumericType; + +/** + * Max value aggregator for star tree + * + * @opensearch.experimental + */ +class MaxValueAggregator extends StatelessDoubleValueAggregator { + + public MaxValueAggregator(StarTreeNumericType starTreeNumericType) { + super(starTreeNumericType, null); + } + + @Override + protected Double performValueAggregation(Double aggregatedValue, Double segmentDocValue) { + return Math.max(aggregatedValue, segmentDocValue); + } +} diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/MetricAggregatorInfo.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/MetricAggregatorInfo.java index a9209a38eca82..8b6db2a183bf8 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/MetricAggregatorInfo.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/MetricAggregatorInfo.java @@ -16,6 +16,7 @@ /** * Builds aggregation function and doc values field pair to support various aggregations + * * @opensearch.experimental */ public class MetricAggregatorInfo implements Comparable { @@ -79,7 +80,15 @@ public StarTreeNumericType getAggregatedValueType() { * @return field name with metric type and field */ public String toFieldName() { - return starFieldName + DELIMITER + field + DELIMITER + metricStat.getTypeName(); + return toFieldName(starFieldName, field, metricStat.getTypeName()); + + } + + /** + * @return field name with star-tree field name metric type and field + */ + public static String toFieldName(String starFieldName, String field, String metricName) { + return starFieldName + DELIMITER + field + DELIMITER + metricName; } @Override @@ -94,7 +103,7 @@ public boolean equals(Object obj) { } if (obj instanceof MetricAggregatorInfo) { MetricAggregatorInfo anotherPair = (MetricAggregatorInfo) obj; - return metricStat == anotherPair.metricStat && field.equals(anotherPair.field); + return metricStat.equals(anotherPair.metricStat) && field.equals(anotherPair.field); } return false; } diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/MinValueAggregator.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/MinValueAggregator.java new file mode 100644 index 0000000000000..46e9188b5dc2f --- /dev/null +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/MinValueAggregator.java @@ -0,0 +1,27 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ +package org.opensearch.index.compositeindex.datacube.startree.aggregators; + +import org.opensearch.index.compositeindex.datacube.startree.aggregators.numerictype.StarTreeNumericType; + +/** + * Min value aggregator for star tree + * + * @opensearch.experimental + */ +class MinValueAggregator extends StatelessDoubleValueAggregator { + + public MinValueAggregator(StarTreeNumericType starTreeNumericType) { + super(starTreeNumericType, null); + } + + @Override + protected Double performValueAggregation(Double aggregatedValue, Double segmentDocValue) { + return Math.min(aggregatedValue, segmentDocValue); + } +} diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/StatelessDoubleValueAggregator.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/StatelessDoubleValueAggregator.java new file mode 100644 index 0000000000000..04c6bcc906ef2 --- /dev/null +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/StatelessDoubleValueAggregator.java @@ -0,0 +1,81 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ +package org.opensearch.index.compositeindex.datacube.startree.aggregators; + +import org.opensearch.index.compositeindex.datacube.startree.aggregators.numerictype.StarTreeNumericType; + +/** + * This is an abstract class that defines the common methods for all double value aggregators + * It is stateless. + * + * @opensearch.experimental + */ +abstract class StatelessDoubleValueAggregator implements ValueAggregator { + + protected final StarTreeNumericType starTreeNumericType; + protected final Double identityValue; + private static final StarTreeNumericType VALUE_AGGREGATOR_TYPE = StarTreeNumericType.DOUBLE; + + public StatelessDoubleValueAggregator(StarTreeNumericType starTreeNumericType, Double identityValue) { + this.starTreeNumericType = starTreeNumericType; + this.identityValue = identityValue; + } + + @Override + public StarTreeNumericType getAggregatedValueType() { + return VALUE_AGGREGATOR_TYPE; + } + + @Override + public Double getInitialAggregatedValueForSegmentDocValue(Long segmentDocValue) { + if (segmentDocValue == null) { + return getIdentityMetricValue(); + } + return starTreeNumericType.getDoubleValue(segmentDocValue); + } + + @Override + public Double mergeAggregatedValues(Double value, Double aggregatedValue) { + if (value == null && aggregatedValue != null) { + return aggregatedValue; + } else if (value != null && aggregatedValue == null) { + return value; + } else if (value == null) { + return getIdentityMetricValue(); + } + return performValueAggregation(value, aggregatedValue); + } + + @Override + public Double toStarTreeNumericTypeValue(Long value) { + try { + if (value == null) { + return getIdentityMetricValue(); + } + return starTreeNumericType.getDoubleValue(value); + } catch (Exception e) { + throw new IllegalStateException("Cannot convert " + value + " to sortable aggregation type", e); + } + } + + @Override + public Double getIdentityMetricValue() { + // the identity value that we return should be inline with the existing aggregations + return identityValue; + } + + /** + * Performs stateless aggregation on the value and the segmentDocValue based on the implementation + * + * @param aggregatedValue aggregated value for the segment so far + * @param segmentDocValue current segment doc value + * @return aggregated value + */ + protected abstract Double performValueAggregation(Double aggregatedValue, Double segmentDocValue); + +} diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/SumValueAggregator.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/SumValueAggregator.java index a471f0e2bd960..1568debd91ae7 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/SumValueAggregator.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/SumValueAggregator.java @@ -7,34 +7,30 @@ */ package org.opensearch.index.compositeindex.datacube.startree.aggregators; -import org.apache.lucene.util.NumericUtils; -import org.opensearch.index.compositeindex.datacube.MetricStat; import org.opensearch.index.compositeindex.datacube.startree.aggregators.numerictype.StarTreeNumericType; import org.opensearch.search.aggregations.metrics.CompensatedSum; /** * Sum value aggregator for star tree * + *

This implementation follows the Kahan summation algorithm to improve the accuracy + * of the sum by tracking and compensating for the accumulated error in each iteration. + * + * @see Kahan Summation Algorithm + * * @opensearch.experimental */ -public class SumValueAggregator implements ValueAggregator { +class SumValueAggregator implements ValueAggregator { - public static final StarTreeNumericType VALUE_AGGREGATOR_TYPE = StarTreeNumericType.DOUBLE; - private double sum = 0; - private double compensation = 0; - private CompensatedSum kahanSummation = new CompensatedSum(0, 0); + private final StarTreeNumericType starTreeNumericType; + private static final StarTreeNumericType VALUE_AGGREGATOR_TYPE = StarTreeNumericType.DOUBLE; - private StarTreeNumericType starTreeNumericType; + private CompensatedSum kahanSummation = new CompensatedSum(0, 0); public SumValueAggregator(StarTreeNumericType starTreeNumericType) { this.starTreeNumericType = starTreeNumericType; } - @Override - public MetricStat getAggregationType() { - return MetricStat.SUM; - } - @Override public StarTreeNumericType getAggregatedValueType() { return VALUE_AGGREGATOR_TYPE; @@ -43,62 +39,60 @@ public StarTreeNumericType getAggregatedValueType() { @Override public Double getInitialAggregatedValueForSegmentDocValue(Long segmentDocValue) { kahanSummation.reset(0, 0); - kahanSummation.add(starTreeNumericType.getDoubleValue(segmentDocValue)); - compensation = kahanSummation.delta(); - sum = kahanSummation.value(); + // add takes care of the sum and compensation internally + if (segmentDocValue != null) { + kahanSummation.add(starTreeNumericType.getDoubleValue(segmentDocValue)); + } else { + kahanSummation.add(getIdentityMetricValue()); + } return kahanSummation.value(); } + // we have overridden this method because the reset with sum and compensation helps us keep + // track of precision and avoids a potential loss in accuracy of sums. @Override public Double mergeAggregatedValueAndSegmentValue(Double value, Long segmentDocValue) { - assert kahanSummation.value() == value; - kahanSummation.reset(sum, compensation); - kahanSummation.add(starTreeNumericType.getDoubleValue(segmentDocValue)); - compensation = kahanSummation.delta(); - sum = kahanSummation.value(); + assert value == null || kahanSummation.value() == value; + // add takes care of the sum and compensation internally + if (segmentDocValue != null) { + kahanSummation.add(starTreeNumericType.getDoubleValue(segmentDocValue)); + } else { + kahanSummation.add(getIdentityMetricValue()); + } return kahanSummation.value(); } @Override public Double mergeAggregatedValues(Double value, Double aggregatedValue) { - assert kahanSummation.value() == aggregatedValue; - kahanSummation.reset(sum, compensation); - kahanSummation.add(value); - compensation = kahanSummation.delta(); - sum = kahanSummation.value(); + assert aggregatedValue == null || kahanSummation.value() == aggregatedValue; + // add takes care of the sum and compensation internally + if (value != null) { + kahanSummation.add(value); + } else { + kahanSummation.add(getIdentityMetricValue()); + } return kahanSummation.value(); } @Override public Double getInitialAggregatedValue(Double value) { kahanSummation.reset(0, 0); - kahanSummation.add(value); - compensation = kahanSummation.delta(); - sum = kahanSummation.value(); - return kahanSummation.value(); - } - - @Override - public int getMaxAggregatedValueByteSize() { - return Double.BYTES; - } - - @Override - public Long toLongValue(Double value) { - try { - return NumericUtils.doubleToSortableLong(value); - } catch (Exception e) { - throw new IllegalStateException("Cannot convert " + value + " to sortable long", e); + // add takes care of the sum and compensation internally + if (value != null) { + kahanSummation.add(value); + } else { + kahanSummation.add(getIdentityMetricValue()); } + return kahanSummation.value(); } @Override public Double toStarTreeNumericTypeValue(Long value) { try { if (value == null) { - return 0.0; + return getIdentityMetricValue(); } - return VALUE_AGGREGATOR_TYPE.getDoubleValue(value); + return starTreeNumericType.getDoubleValue(value); } catch (Exception e) { throw new IllegalStateException("Cannot convert " + value + " to sortable aggregation type", e); } @@ -106,6 +100,7 @@ public Double toStarTreeNumericTypeValue(Long value) { @Override public Double getIdentityMetricValue() { + // in present aggregations, if the metric behind sum is missing, we treat it as 0 return 0D; } } diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/ValueAggregator.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/ValueAggregator.java index 048582cc530e5..39bf235a83409 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/ValueAggregator.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/ValueAggregator.java @@ -7,7 +7,6 @@ */ package org.opensearch.index.compositeindex.datacube.startree.aggregators; -import org.opensearch.index.compositeindex.datacube.MetricStat; import org.opensearch.index.compositeindex.datacube.startree.aggregators.numerictype.StarTreeNumericType; /** @@ -17,11 +16,6 @@ */ public interface ValueAggregator { - /** - * Returns the type of the aggregation. - */ - MetricStat getAggregationType(); - /** * Returns the data type of the aggregated value. */ @@ -35,7 +29,10 @@ public interface ValueAggregator { /** * Applies a segment doc value to the current aggregated value. */ - A mergeAggregatedValueAndSegmentValue(A value, Long segmentDocValue); + default A mergeAggregatedValueAndSegmentValue(A value, Long segmentDocValue) { + A aggregatedValue = getInitialAggregatedValueForSegmentDocValue(segmentDocValue); + return mergeAggregatedValues(value, aggregatedValue); + } /** * Applies an aggregated value to the current aggregated value. @@ -45,17 +42,12 @@ public interface ValueAggregator { /** * Clones an aggregated value. */ - A getInitialAggregatedValue(A value); - - /** - * Returns the maximum size in bytes of the aggregated values seen so far. - */ - int getMaxAggregatedValueByteSize(); - - /** - * Converts an aggregated value into a Long type. - */ - Long toLongValue(A value); + default A getInitialAggregatedValue(A value) { + if (value == null) { + return getIdentityMetricValue(); + } + return value; + } /** * Converts an aggregated value from a Long type. diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/ValueAggregatorFactory.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/ValueAggregatorFactory.java index 240bbd37a53ee..5e071e2491d19 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/ValueAggregatorFactory.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/ValueAggregatorFactory.java @@ -27,31 +27,18 @@ private ValueAggregatorFactory() {} */ public static ValueAggregator getValueAggregator(MetricStat aggregationType, StarTreeNumericType starTreeNumericType) { switch (aggregationType) { - // other metric types (count, min, max, avg) will be supported in the future + // avg aggregator will be covered in the part of query (using count and sum) case SUM: return new SumValueAggregator(starTreeNumericType); case COUNT: return new CountValueAggregator(starTreeNumericType); + case MIN: + return new MinValueAggregator(starTreeNumericType); + case MAX: + return new MaxValueAggregator(starTreeNumericType); default: throw new IllegalStateException("Unsupported aggregation type: " + aggregationType); } } - /** - * Returns the data type of the aggregated value for the given aggregation type. - * - * @param aggregationType Aggregation type - * @return Data type of the aggregated value - */ - public static StarTreeNumericType getAggregatedValueType(MetricStat aggregationType) { - switch (aggregationType) { - // other metric types (count, min, max, avg) will be supported in the future - case SUM: - return SumValueAggregator.VALUE_AGGREGATOR_TYPE; - case COUNT: - return CountValueAggregator.VALUE_AGGREGATOR_TYPE; - default: - throw new IllegalStateException("Unsupported aggregation type: " + aggregationType); - } - } } diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/AbstractDocumentsFileManager.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/AbstractDocumentsFileManager.java index 78c49dbada6b2..4214a46b2fc1c 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/AbstractDocumentsFileManager.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/AbstractDocumentsFileManager.java @@ -175,12 +175,7 @@ protected long readMetrics(RandomAccessInput input, long offset, int numMetrics, throw new IllegalStateException("Unsupported metric type"); } } - offset += StarTreeDocumentBitSetUtil.readBitSet( - input, - offset, - metrics, - index -> metricAggregatorInfos.get(index).getValueAggregators().getIdentityMetricValue() - ); + offset += StarTreeDocumentBitSetUtil.readBitSet(input, offset, metrics, index -> null); return offset; } diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/BaseStarTreeBuilder.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/BaseStarTreeBuilder.java index 56bb46e83a9da..872826aa6db06 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/BaseStarTreeBuilder.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/BaseStarTreeBuilder.java @@ -79,9 +79,9 @@ public abstract class BaseStarTreeBuilder implements StarTreeBuilder { /** * Reads all the configuration related to dimensions and metrics, builds a star-tree based on the different construction parameters. * - * @param starTreeField holds the configuration for the star tree - * @param state stores the segment write state - * @param mapperService helps to find the original type of the field + * @param starTreeField holds the configuration for the star tree + * @param state stores the segment write state + * @param mapperService helps to find the original type of the field */ protected BaseStarTreeBuilder(StarTreeField starTreeField, SegmentWriteState state, MapperService mapperService) { logger.debug("Building star tree : {}", starTreeField.getName()); @@ -208,7 +208,7 @@ protected StarTreeDocument getStarTreeDocument( * aggregated star-tree documents. * * @param dimensionReaders List of docValues readers to read dimensions from the segment - * @param metricReaders List of docValues readers to read metrics from the segment + * @param metricReaders List of docValues readers to read metrics from the segment * @return Iterator for the aggregated star-tree document */ public abstract Iterator sortAndAggregateSegmentDocuments( @@ -229,7 +229,6 @@ public abstract Iterator generateStarTreeDocumentsForStarNode( /** * Returns the star-tree document from the segment based on the current doc id - * */ protected StarTreeDocument getSegmentStarTreeDocument( int currentDocId, @@ -354,32 +353,15 @@ protected StarTreeDocument reduceSegmentStarTreeDocuments( /** * Safely converts the metric value of object type to long. + * Nulls are handled during aggregation * * @param metric value of the metric * @return converted metric value to long */ - private static long getLong(Object metric) { + private static Long getLong(Object metric) { Long metricValue = null; - // TODO : remove this after we merge identity changes - if (metric instanceof Double) { - if (0D == (double) metric) { - return 0L; - } - } - try { - if (metric instanceof Long) { - metricValue = (long) metric; - } else if (metric != null) { - metricValue = Long.valueOf(String.valueOf(metric)); - } - } catch (Exception e) { - throw new IllegalStateException("unable to cast segment metric", e); - } - - if (metricValue == null) { - return 0; - // TODO: handle this properly - // throw new IllegalStateException("unable to cast segment metric"); + if (metric instanceof Long) { + metricValue = (long) metric; } return metricValue; } @@ -428,7 +410,6 @@ public StarTreeDocument reduceStarTreeDocuments(StarTreeDocument aggregatedDocum * Builds the star tree from the original segment documents * * @param fieldProducerMap contain s the docValues producer to get docValues associated with each field - * * @throws IOException when we are unable to build star-tree */ public void build(Map fieldProducerMap) throws IOException { @@ -493,14 +474,10 @@ public List getMetricReaders(SegmentWriteState stat if (metricFieldInfo == null) { metricFieldInfo = getFieldInfo(metric.getField()); } - // TODO - // if (metricStat != MetricStat.COUNT) { - // Need not initialize the metric reader for COUNT metric type + SequentialDocValuesIterator metricReader = new SequentialDocValuesIterator( fieldProducerMap.get(metricFieldInfo.name).getSortedNumeric(metricFieldInfo) ); - // } - metricReaders.add(metricReader); } } diff --git a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/AbstractValueAggregatorTests.java b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/AbstractValueAggregatorTests.java new file mode 100644 index 0000000000000..f6adf442bb6ab --- /dev/null +++ b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/AbstractValueAggregatorTests.java @@ -0,0 +1,77 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.index.compositeindex.datacube.startree.aggregators; + +import com.carrotsearch.randomizedtesting.annotations.ParametersFactory; + +import org.opensearch.index.compositeindex.datacube.startree.aggregators.numerictype.StarTreeNumericType; +import org.opensearch.test.OpenSearchTestCase; +import org.junit.Before; + +import java.util.ArrayList; +import java.util.Collection; +import java.util.List; + +public abstract class AbstractValueAggregatorTests extends OpenSearchTestCase { + + private ValueAggregator aggregator; + private StarTreeNumericType starTreeNumericType; + + public AbstractValueAggregatorTests(StarTreeNumericType starTreeNumericType) { + this.starTreeNumericType = starTreeNumericType; + } + + @Before + public void setup() { + aggregator = getValueAggregator(starTreeNumericType); + } + + @ParametersFactory + public static Collection parameters() { + List parameters = new ArrayList<>(); + for (StarTreeNumericType starTreeNumericType : StarTreeNumericType.values()) { + parameters.add(new Object[] { starTreeNumericType }); + } + return parameters; + } + + public abstract ValueAggregator getValueAggregator(StarTreeNumericType starTreeNumericType); + + public void testGetInitialAggregatedValueForSegmentDocNullValue() { + assertEquals(aggregator.getIdentityMetricValue(), aggregator.getInitialAggregatedValueForSegmentDocValue(null)); + } + + public void testMergeAggregatedNullValueAndSegmentNullValue() { + if (aggregator instanceof CountValueAggregator) { + assertThrows(AssertionError.class, () -> aggregator.mergeAggregatedValueAndSegmentValue(null, null)); + } else { + assertEquals(aggregator.getIdentityMetricValue(), aggregator.mergeAggregatedValueAndSegmentValue(null, null)); + } + } + + public void testMergeAggregatedNullValues() { + assertEquals(aggregator.getIdentityMetricValue(), aggregator.mergeAggregatedValues(null, null)); + } + + public void testGetInitialAggregatedNullValue() { + assertEquals(aggregator.getIdentityMetricValue(), aggregator.getInitialAggregatedValue(null)); + } + + public void testGetInitialAggregatedValueForSegmentDocValue() { + long randomLong = randomLong(); + if (aggregator instanceof CountValueAggregator) { + assertEquals(CountValueAggregator.DEFAULT_INITIAL_VALUE, aggregator.getInitialAggregatedValueForSegmentDocValue(randomLong())); + } else { + assertEquals( + aggregator.toStarTreeNumericTypeValue(randomLong), + aggregator.getInitialAggregatedValueForSegmentDocValue(randomLong) + ); + } + } +} diff --git a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/CountValueAggregatorTests.java b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/CountValueAggregatorTests.java index 8e6e9e9974646..7389d68987898 100644 --- a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/CountValueAggregatorTests.java +++ b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/CountValueAggregatorTests.java @@ -8,46 +8,48 @@ package org.opensearch.index.compositeindex.datacube.startree.aggregators; -import org.opensearch.index.compositeindex.datacube.MetricStat; import org.opensearch.index.compositeindex.datacube.startree.aggregators.numerictype.StarTreeNumericType; -import org.opensearch.test.OpenSearchTestCase; -public class CountValueAggregatorTests extends OpenSearchTestCase { - private final CountValueAggregator aggregator = new CountValueAggregator(StarTreeNumericType.LONG); +public class CountValueAggregatorTests extends AbstractValueAggregatorTests { - public void testGetAggregationType() { - assertEquals(MetricStat.COUNT.getTypeName(), aggregator.getAggregationType().getTypeName()); - } - - public void testGetAggregatedValueType() { - assertEquals(CountValueAggregator.VALUE_AGGREGATOR_TYPE, aggregator.getAggregatedValueType()); - } + private CountValueAggregator aggregator; - public void testGetInitialAggregatedValueForSegmentDocValue() { - assertEquals(1L, aggregator.getInitialAggregatedValueForSegmentDocValue(randomLong()), 0.0); + public CountValueAggregatorTests(StarTreeNumericType starTreeNumericType) { + super(starTreeNumericType); } public void testMergeAggregatedValueAndSegmentValue() { - assertEquals(3L, aggregator.mergeAggregatedValueAndSegmentValue(2L, 3L), 0.0); + long randomLong = randomLong(); + assertEquals(randomLong + 1, aggregator.mergeAggregatedValueAndSegmentValue(randomLong, 3L), 0.0); } public void testMergeAggregatedValues() { - assertEquals(5L, aggregator.mergeAggregatedValues(2L, 3L), 0.0); + long randomLong1 = randomLong(); + long randomLong2 = randomLong(); + assertEquals(randomLong1 + randomLong2, aggregator.mergeAggregatedValues(randomLong1, randomLong2), 0.0); + assertEquals(randomLong1, aggregator.mergeAggregatedValues(randomLong1, null), 0.0); + assertEquals(randomLong2, aggregator.mergeAggregatedValues(null, randomLong2), 0.0); } public void testGetInitialAggregatedValue() { - assertEquals(3L, aggregator.getInitialAggregatedValue(3L), 0.0); + long randomLong = randomLong(); + assertEquals(randomLong, aggregator.getInitialAggregatedValue(randomLong), 0.0); } - public void testGetMaxAggregatedValueByteSize() { - assertEquals(Long.BYTES, aggregator.getMaxAggregatedValueByteSize()); + public void testToStarTreeNumericTypeValue() { + long randomLong = randomLong(); + assertEquals(randomLong, aggregator.toStarTreeNumericTypeValue(randomLong), 0.0); + assertNull(aggregator.toStarTreeNumericTypeValue(null)); } - public void testToLongValue() { - assertEquals(3L, aggregator.toLongValue(3L), 0.0); + public void testIdentityMetricValue() { + assertEquals(0L, aggregator.getIdentityMetricValue(), 0); } - public void testToStarTreeNumericTypeValue() { - assertEquals(3L, aggregator.toStarTreeNumericTypeValue(3L), 0.0); + @Override + public ValueAggregator getValueAggregator(StarTreeNumericType starTreeNumericType) { + aggregator = new CountValueAggregator(starTreeNumericType); + return aggregator; } + } diff --git a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/MaxValueAggregatorTests.java b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/MaxValueAggregatorTests.java new file mode 100644 index 0000000000000..a4e01bb70492c --- /dev/null +++ b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/MaxValueAggregatorTests.java @@ -0,0 +1,72 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.index.compositeindex.datacube.startree.aggregators; + +import org.apache.lucene.util.NumericUtils; +import org.opensearch.index.compositeindex.datacube.startree.aggregators.numerictype.StarTreeNumericType; + +public class MaxValueAggregatorTests extends AbstractValueAggregatorTests { + + private MaxValueAggregator aggregator; + + public MaxValueAggregatorTests(StarTreeNumericType starTreeNumericType) { + super(starTreeNumericType); + } + + public void testMergeAggregatedValueAndSegmentValue() { + Long randomLong = randomLong(); + double randomDouble = randomDouble(); + assertEquals( + Math.max(aggregator.toStarTreeNumericTypeValue(randomLong), randomDouble), + aggregator.mergeAggregatedValueAndSegmentValue(randomDouble, randomLong), + 0.0 + ); + assertEquals( + aggregator.toStarTreeNumericTypeValue(randomLong), + aggregator.mergeAggregatedValueAndSegmentValue(null, randomLong), + 0.0 + ); + assertEquals(randomDouble, aggregator.mergeAggregatedValueAndSegmentValue(randomDouble, null), 0.0); + assertEquals( + Math.max(2.0, aggregator.toStarTreeNumericTypeValue(3L)), + aggregator.mergeAggregatedValueAndSegmentValue(2.0, 3L), + 0.0 + ); + } + + public void testMergeAggregatedValues() { + double randomDouble = randomDouble(); + assertEquals(randomDouble, aggregator.mergeAggregatedValues(Double.MIN_VALUE, randomDouble), 0.0); + assertEquals(randomDouble, aggregator.mergeAggregatedValues(null, randomDouble), 0.0); + assertEquals(randomDouble, aggregator.mergeAggregatedValues(randomDouble, null), 0.0); + assertEquals(3.0, aggregator.mergeAggregatedValues(2.0, 3.0), 0.0); + } + + public void testGetInitialAggregatedValue() { + double randomDouble = randomDouble(); + assertEquals(randomDouble, aggregator.getInitialAggregatedValue(randomDouble), 0.0); + } + + public void testToStarTreeNumericTypeValue() { + MaxValueAggregator aggregator = new MaxValueAggregator(StarTreeNumericType.DOUBLE); + long randomLong = randomLong(); + assertEquals(NumericUtils.sortableLongToDouble(randomLong), aggregator.toStarTreeNumericTypeValue(randomLong), 0.0); + } + + public void testIdentityMetricValue() { + assertNull(aggregator.getIdentityMetricValue()); + } + + @Override + public ValueAggregator getValueAggregator(StarTreeNumericType starTreeNumericType) { + aggregator = new MaxValueAggregator(starTreeNumericType); + return aggregator; + } + +} diff --git a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/MinValueAggregatorTests.java b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/MinValueAggregatorTests.java new file mode 100644 index 0000000000000..b8233b0fd7fe0 --- /dev/null +++ b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/MinValueAggregatorTests.java @@ -0,0 +1,71 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.index.compositeindex.datacube.startree.aggregators; + +import org.apache.lucene.util.NumericUtils; +import org.opensearch.index.compositeindex.datacube.startree.aggregators.numerictype.StarTreeNumericType; + +public class MinValueAggregatorTests extends AbstractValueAggregatorTests { + private MinValueAggregator aggregator; + + public MinValueAggregatorTests(StarTreeNumericType starTreeNumericType) { + super(starTreeNumericType); + } + + public void testMergeAggregatedValueAndSegmentValue() { + Long randomLong = randomLong(); + double randomDouble = randomDouble(); + assertEquals( + Math.min(aggregator.toStarTreeNumericTypeValue(randomLong), randomDouble), + aggregator.mergeAggregatedValueAndSegmentValue(randomDouble, randomLong), + 0.0 + ); + assertEquals( + aggregator.toStarTreeNumericTypeValue(randomLong), + aggregator.mergeAggregatedValueAndSegmentValue(null, randomLong), + 0.0 + ); + assertEquals(randomDouble, aggregator.mergeAggregatedValueAndSegmentValue(randomDouble, null), 0.0); + assertEquals( + Math.min(2.0, aggregator.toStarTreeNumericTypeValue(3L)), + aggregator.mergeAggregatedValueAndSegmentValue(2.0, 3L), + 0.0 + ); + } + + public void testMergeAggregatedValues() { + double randomDouble = randomDouble(); + assertEquals(randomDouble, aggregator.mergeAggregatedValues(Double.MAX_VALUE, randomDouble), 0.0); + assertEquals(randomDouble, aggregator.mergeAggregatedValues(null, randomDouble), 0.0); + assertEquals(randomDouble, aggregator.mergeAggregatedValues(randomDouble, null), 0.0); + assertEquals(2.0, aggregator.mergeAggregatedValues(2.0, 3.0), 0.0); + } + + public void testGetInitialAggregatedValue() { + double randomDouble = randomDouble(); + assertEquals(randomDouble, aggregator.getInitialAggregatedValue(randomDouble), 0.0); + } + + public void testToStarTreeNumericTypeValue() { + MinValueAggregator aggregator = new MinValueAggregator(StarTreeNumericType.DOUBLE); + long randomLong = randomLong(); + assertEquals(NumericUtils.sortableLongToDouble(randomLong), aggregator.toStarTreeNumericTypeValue(randomLong), 0.0); + } + + public void testIdentityMetricValue() { + assertNull(aggregator.getIdentityMetricValue()); + } + + @Override + public ValueAggregator getValueAggregator(StarTreeNumericType starTreeNumericType) { + aggregator = new MinValueAggregator(starTreeNumericType); + return aggregator; + } + +} diff --git a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/StaticValueAggregatorTests.java b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/StaticValueAggregatorTests.java new file mode 100644 index 0000000000000..487bd54602503 --- /dev/null +++ b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/StaticValueAggregatorTests.java @@ -0,0 +1,133 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.index.compositeindex.datacube.startree.aggregators; + +import org.apache.lucene.util.NumericUtils; +import org.opensearch.index.compositeindex.datacube.startree.aggregators.numerictype.StarTreeNumericType; +import org.opensearch.search.aggregations.metrics.CompensatedSum; +import org.opensearch.test.OpenSearchTestCase; + +public class StaticValueAggregatorTests extends OpenSearchTestCase { + + // tests the extreme case where normal sum will lose precision + public void testKahanSummation() { + double[] numbers = { 1e-16, 1, -1e-16 }; + double expected = 1; + + // initializing our sum aggregator to derive exact sum using kahan summation + double aggregatedValue = getAggregatedValue(numbers); + assertEquals(expected, aggregatedValue, 0); + + // assert kahan summation plain logic with our aggregated value + double actual = kahanSum(numbers); + assertEquals(actual, aggregatedValue, 0); + + // assert that normal sum fails for this case + double normalSum = normalSum(numbers); + assertNotEquals(expected, normalSum, 0); + assertNotEquals(actual, normalSum, 0); + assertNotEquals(aggregatedValue, normalSum, 0); + + } + + private static double getAggregatedValue(double[] numbers) { + // explicitly took double to test for most precision + // hard to run similar tests for different data types dynamically as inputs and precision vary + SumValueAggregator aggregator = new SumValueAggregator(StarTreeNumericType.DOUBLE); + double aggregatedValue = aggregator.getInitialAggregatedValueForSegmentDocValue(NumericUtils.doubleToSortableLong(numbers[0])); + aggregatedValue = aggregator.mergeAggregatedValueAndSegmentValue(aggregatedValue, NumericUtils.doubleToSortableLong(numbers[1])); + aggregatedValue = aggregator.mergeAggregatedValueAndSegmentValue(aggregatedValue, NumericUtils.doubleToSortableLong(numbers[2])); + return aggregatedValue; + } + + private double kahanSum(double[] numbers) { + CompensatedSum compensatedSum = new CompensatedSum(0, 0); + for (double num : numbers) { + compensatedSum.add(num); + } + return compensatedSum.value(); + } + + private double normalSum(double[] numbers) { + double sum = 0.0; + for (double num : numbers) { + sum += num; + } + return sum; + } + + public void testMaxAggregatorExtremeValues() { + double[] numbers = { Double.MAX_VALUE, Double.MIN_VALUE, 0.0, Double.MAX_VALUE + 1 }; + double expected = Double.MAX_VALUE + 1; + MaxValueAggregator aggregator = new MaxValueAggregator(StarTreeNumericType.DOUBLE); + double aggregatedValue = aggregator.getInitialAggregatedValueForSegmentDocValue(NumericUtils.doubleToSortableLong(numbers[0])); + for (int i = 1; i < numbers.length; i++) { + aggregatedValue = aggregator.mergeAggregatedValueAndSegmentValue( + aggregatedValue, + NumericUtils.doubleToSortableLong(numbers[i]) + ); + } + assertEquals(expected, aggregatedValue, 0); + } + + public void testMaxAggregatorExtremeValues_Infinity() { + double[] numbers = { + Double.MAX_VALUE, + Double.MIN_VALUE, + 0.0, + Double.MAX_VALUE + 1, + Double.POSITIVE_INFINITY, + Double.NEGATIVE_INFINITY }; + double expected = Double.POSITIVE_INFINITY; + MaxValueAggregator aggregator = new MaxValueAggregator(StarTreeNumericType.DOUBLE); + double aggregatedValue = aggregator.getInitialAggregatedValueForSegmentDocValue(NumericUtils.doubleToSortableLong(numbers[0])); + for (int i = 1; i < numbers.length; i++) { + aggregatedValue = aggregator.mergeAggregatedValueAndSegmentValue( + aggregatedValue, + NumericUtils.doubleToSortableLong(numbers[i]) + ); + } + assertEquals(expected, aggregatedValue, 0); + } + + public void testMinAggregatorExtremeValues() { + double[] numbers = { Double.MAX_VALUE, Double.MIN_VALUE - 1, 0.0, Double.MAX_VALUE + 1 }; + double expected = Double.MIN_VALUE - 1; + MinValueAggregator aggregator = new MinValueAggregator(StarTreeNumericType.DOUBLE); + double aggregatedValue = aggregator.getInitialAggregatedValueForSegmentDocValue(NumericUtils.doubleToSortableLong(numbers[0])); + for (int i = 1; i < numbers.length; i++) { + aggregatedValue = aggregator.mergeAggregatedValueAndSegmentValue( + aggregatedValue, + NumericUtils.doubleToSortableLong(numbers[i]) + ); + } + assertEquals(expected, aggregatedValue, 0); + } + + public void testMinAggregatorExtremeValues_Infinity() { + double[] numbers = { + Double.MAX_VALUE, + Double.MIN_VALUE, + 0.0, + Double.MAX_VALUE + 1, + Double.POSITIVE_INFINITY, + Double.NEGATIVE_INFINITY }; + double expected = Double.NEGATIVE_INFINITY; + MinValueAggregator aggregator = new MinValueAggregator(StarTreeNumericType.DOUBLE); + double aggregatedValue = aggregator.getInitialAggregatedValueForSegmentDocValue(NumericUtils.doubleToSortableLong(numbers[0])); + for (int i = 1; i < numbers.length; i++) { + aggregatedValue = aggregator.mergeAggregatedValueAndSegmentValue( + aggregatedValue, + NumericUtils.doubleToSortableLong(numbers[i]) + ); + } + assertEquals(expected, aggregatedValue, 0); + } + +} diff --git a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/SumValueAggregatorTests.java b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/SumValueAggregatorTests.java index dd66d4344c9e8..3a85357da5ffe 100644 --- a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/SumValueAggregatorTests.java +++ b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/SumValueAggregatorTests.java @@ -8,62 +8,75 @@ package org.opensearch.index.compositeindex.datacube.startree.aggregators; -import org.apache.lucene.util.NumericUtils; -import org.opensearch.index.compositeindex.datacube.MetricStat; import org.opensearch.index.compositeindex.datacube.startree.aggregators.numerictype.StarTreeNumericType; -import org.opensearch.test.OpenSearchTestCase; -import org.junit.Before; -public class SumValueAggregatorTests extends OpenSearchTestCase { +public class SumValueAggregatorTests extends AbstractValueAggregatorTests { private SumValueAggregator aggregator; - @Before - public void setup() { - aggregator = new SumValueAggregator(StarTreeNumericType.LONG); + public SumValueAggregatorTests(StarTreeNumericType starTreeNumericType) { + super(starTreeNumericType); } - public void testGetAggregationType() { - assertEquals(MetricStat.SUM.getTypeName(), aggregator.getAggregationType().getTypeName()); - } - - public void testGetAggregatedValueType() { - assertEquals(SumValueAggregator.VALUE_AGGREGATOR_TYPE, aggregator.getAggregatedValueType()); - } - - public void testGetInitialAggregatedValueForSegmentDocValue() { - assertEquals(1.0, aggregator.getInitialAggregatedValueForSegmentDocValue(1L), 0.0); - assertThrows(NullPointerException.class, () -> aggregator.getInitialAggregatedValueForSegmentDocValue(null)); + @Override + public ValueAggregator getValueAggregator(StarTreeNumericType starTreeNumericType) { + aggregator = new SumValueAggregator(starTreeNumericType); + return aggregator; } public void testMergeAggregatedValueAndSegmentValue() { - aggregator.getInitialAggregatedValue(2.0); - assertEquals(5.0, aggregator.mergeAggregatedValueAndSegmentValue(2.0, 3L), 0.0); + double randomDouble = randomDouble(); + Long randomLong = randomLong(); + aggregator.getInitialAggregatedValue(randomDouble); + assertEquals( + randomDouble + aggregator.toStarTreeNumericTypeValue(randomLong), + aggregator.mergeAggregatedValueAndSegmentValue(randomDouble, randomLong), + 0.0 + ); } public void testMergeAggregatedValueAndSegmentValue_nullSegmentDocValue() { - aggregator.getInitialAggregatedValue(2.0); - assertThrows(NullPointerException.class, () -> aggregator.mergeAggregatedValueAndSegmentValue(2.0, null)); + double randomDouble1 = randomDouble(); + Long randomLong = randomLong(); + aggregator.getInitialAggregatedValue(randomDouble1); + assertEquals(randomDouble1, aggregator.mergeAggregatedValueAndSegmentValue(randomDouble1, null), 0.0); + assertEquals( + randomDouble1 + aggregator.toStarTreeNumericTypeValue(randomLong), + aggregator.mergeAggregatedValueAndSegmentValue(randomDouble1, randomLong), + 0.0 + ); + } + + public void testMergeAggregatedValueAndSegmentValue_nullInitialDocValue() { + Long randomLong = randomLong(); + aggregator.getInitialAggregatedValue(null); + assertEquals( + aggregator.toStarTreeNumericTypeValue(randomLong), + aggregator.mergeAggregatedValueAndSegmentValue(null, randomLong), + 0.0 + ); } public void testMergeAggregatedValues() { - aggregator.getInitialAggregatedValue(3.0); - assertEquals(5.0, aggregator.mergeAggregatedValues(2.0, 3.0), 0.0); + double randomDouble1 = randomDouble(); + double randomDouble2 = randomDouble(); + aggregator.getInitialAggregatedValue(randomDouble1); + assertEquals(randomDouble1, aggregator.mergeAggregatedValues(null, randomDouble1), 0.0); + assertEquals(randomDouble1 + randomDouble2, aggregator.mergeAggregatedValues(randomDouble2, randomDouble1), 0.0); } public void testGetInitialAggregatedValue() { - assertEquals(3.14, aggregator.getInitialAggregatedValue(3.14), 0.0); + double randomDouble = randomDouble(); + assertEquals(randomDouble, aggregator.getInitialAggregatedValue(randomDouble), 0.0); } - public void testGetMaxAggregatedValueByteSize() { - assertEquals(Double.BYTES, aggregator.getMaxAggregatedValueByteSize()); + public void testToStarTreeNumericTypeValue() { + long randomLong = randomLong(); + assertEquals(aggregator.toStarTreeNumericTypeValue(randomLong), aggregator.toStarTreeNumericTypeValue(randomLong), 0.0); } - public void testToLongValue() { - assertEquals(NumericUtils.doubleToSortableLong(3.14), aggregator.toLongValue(3.14), 0.0); + public void testIdentityMetricValue() { + assertEquals(0.0, aggregator.getIdentityMetricValue(), 0); } - public void testToStarTreeNumericTypeValue() { - assertEquals(NumericUtils.sortableLongToDouble(3L), aggregator.toStarTreeNumericTypeValue(3L), 0.0); - } } diff --git a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/ValueAggregatorFactoryTests.java b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/ValueAggregatorFactoryTests.java index 428668511fb2e..5e0bedf5e06a5 100644 --- a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/ValueAggregatorFactoryTests.java +++ b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/ValueAggregatorFactoryTests.java @@ -20,8 +20,29 @@ public void testGetValueAggregatorForSumType() { assertEquals(SumValueAggregator.class, aggregator.getClass()); } - public void testGetAggregatedValueTypeForSumType() { - StarTreeNumericType starTreeNumericType = ValueAggregatorFactory.getAggregatedValueType(MetricStat.SUM); - assertEquals(SumValueAggregator.VALUE_AGGREGATOR_TYPE, starTreeNumericType); + public void testGetValueAggregatorForMinType() { + ValueAggregator aggregator = ValueAggregatorFactory.getValueAggregator(MetricStat.MIN, StarTreeNumericType.LONG); + assertNotNull(aggregator); + assertEquals(MinValueAggregator.class, aggregator.getClass()); + } + + public void testGetValueAggregatorForMaxType() { + ValueAggregator aggregator = ValueAggregatorFactory.getValueAggregator(MetricStat.MAX, StarTreeNumericType.LONG); + assertNotNull(aggregator); + assertEquals(MaxValueAggregator.class, aggregator.getClass()); + } + + public void testGetValueAggregatorForCountType() { + ValueAggregator aggregator = ValueAggregatorFactory.getValueAggregator(MetricStat.COUNT, StarTreeNumericType.LONG); + assertNotNull(aggregator); + assertEquals(CountValueAggregator.class, aggregator.getClass()); } + + public void testGetValueAggregatorForAvgType() { + assertThrows( + IllegalStateException.class, + () -> ValueAggregatorFactory.getValueAggregator(MetricStat.AVG, StarTreeNumericType.LONG) + ); + } + } diff --git a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/AbstractStarTreeBuilderTests.java b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/AbstractStarTreeBuilderTests.java index 131d7444ff91c..d1a85949da7fe 100644 --- a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/AbstractStarTreeBuilderTests.java +++ b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/AbstractStarTreeBuilderTests.java @@ -92,7 +92,9 @@ public void setup() throws IOException { metrics = List.of( new Metric("field2", List.of(MetricStat.SUM)), new Metric("field4", List.of(MetricStat.SUM)), - new Metric("field6", List.of(MetricStat.COUNT)) + new Metric("field6", List.of(MetricStat.COUNT)), + new Metric("field9", List.of(MetricStat.MIN)), + new Metric("field10", List.of(MetricStat.MAX)) ); DocValuesProducer docValuesProducer = mock(DocValuesProducer.class); @@ -141,8 +143,12 @@ public void setup() throws IOException { .build(new Mapper.BuilderContext(settings, new ContentPath())); NumberFieldMapper numberFieldMapper3 = new NumberFieldMapper.Builder("field6", NumberFieldMapper.NumberType.DOUBLE, false, true) .build(new Mapper.BuilderContext(settings, new ContentPath())); + NumberFieldMapper numberFieldMapper4 = new NumberFieldMapper.Builder("field9", NumberFieldMapper.NumberType.DOUBLE, false, true) + .build(new Mapper.BuilderContext(settings, new ContentPath())); + NumberFieldMapper numberFieldMapper5 = new NumberFieldMapper.Builder("field10", NumberFieldMapper.NumberType.DOUBLE, false, true) + .build(new Mapper.BuilderContext(settings, new ContentPath())); MappingLookup fieldMappers = new MappingLookup( - Set.of(numberFieldMapper1, numberFieldMapper2, numberFieldMapper3), + Set.of(numberFieldMapper1, numberFieldMapper2, numberFieldMapper3, numberFieldMapper4, numberFieldMapper5), Collections.emptyList(), Collections.emptyList(), 0, @@ -181,22 +187,27 @@ public void test_sortAndAggregateStarTreeDocuments() throws IOException { int noOfStarTreeDocuments = 5; StarTreeDocument[] starTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; - starTreeDocuments[0] = new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Double[] { 12.0, 10.0, randomDouble() }); - starTreeDocuments[1] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Double[] { 10.0, 6.0, randomDouble() }); - starTreeDocuments[2] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Double[] { 14.0, 12.0, randomDouble() }); - starTreeDocuments[3] = new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Double[] { 9.0, 4.0, randomDouble() }); - starTreeDocuments[4] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Double[] { 11.0, 16.0, randomDouble() }); + starTreeDocuments[0] = new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Double[] { 12.0, 10.0, randomDouble(), 8.0, 20.0 }); + starTreeDocuments[1] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Double[] { 10.0, 6.0, randomDouble(), 12.0, 10.0 }); + starTreeDocuments[2] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Double[] { 14.0, 12.0, randomDouble(), 6.0, 24.0 }); + starTreeDocuments[3] = new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Double[] { 9.0, 4.0, randomDouble(), 9.0, 12.0 }); + starTreeDocuments[4] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Double[] { 11.0, 16.0, randomDouble(), 8.0, 13.0 }); StarTreeDocument[] segmentStarTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; for (int i = 0; i < noOfStarTreeDocuments; i++) { long metric1 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[0]); long metric2 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[1]); long metric3 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[2]); - segmentStarTreeDocuments[i] = new StarTreeDocument(starTreeDocuments[i].dimensions, new Long[] { metric1, metric2, metric3 }); + long metric4 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[3]); + long metric5 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[4]); + segmentStarTreeDocuments[i] = new StarTreeDocument( + starTreeDocuments[i].dimensions, + new Long[] { metric1, metric2, metric3, metric4, metric5 } + ); } List inorderStarTreeDocuments = List.of( - new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Object[] { 21.0, 14.0, 2L }), - new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Object[] { 35.0, 34.0, 3L }) + new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Object[] { 21.0, 14.0, 2L, 8.0, 20.0 }), + new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Object[] { 35.0, 34.0, 3L, 6.0, 24.0 }) ); Iterator expectedStarTreeDocumentIterator = inorderStarTreeDocuments.iterator(); @@ -220,6 +231,8 @@ public void test_sortAndAggregateStarTreeDocuments() throws IOException { assertEquals(expectedStarTreeDocument.metrics[0], resultStarTreeDocument.metrics[0]); assertEquals(expectedStarTreeDocument.metrics[1], resultStarTreeDocument.metrics[1]); assertEquals(expectedStarTreeDocument.metrics[2], resultStarTreeDocument.metrics[2]); + assertEquals(expectedStarTreeDocument.metrics[3], resultStarTreeDocument.metrics[3]); + assertEquals(expectedStarTreeDocument.metrics[4], resultStarTreeDocument.metrics[4]); numOfAggregatedDocuments++; } @@ -267,26 +280,31 @@ public void test_sortAndAggregateStarTreeDocuments_nullMetric() throws IOExcepti int noOfStarTreeDocuments = 5; StarTreeDocument[] starTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; - starTreeDocuments[0] = new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Double[] { 12.0, 10.0, randomDouble() }); - starTreeDocuments[1] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Double[] { 10.0, 6.0, randomDouble() }); - starTreeDocuments[2] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Double[] { 14.0, 12.0, randomDouble() }); - starTreeDocuments[3] = new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Double[] { 9.0, 4.0, randomDouble() }); - starTreeDocuments[4] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Double[] { 11.0, null, randomDouble() }); + starTreeDocuments[0] = new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Double[] { 12.0, 10.0, randomDouble(), 8.0, 20.0 }); + starTreeDocuments[1] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Double[] { 10.0, 6.0, randomDouble(), 12.0, 10.0 }); + starTreeDocuments[2] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Double[] { 14.0, 12.0, randomDouble(), 6.0, 24.0 }); + starTreeDocuments[3] = new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Double[] { 9.0, 4.0, randomDouble(), 9.0, 12.0 }); + starTreeDocuments[4] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Double[] { 11.0, null, randomDouble(), 8.0, 13.0 }); List inorderStarTreeDocuments = List.of( - new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Object[] { 21.0, 14.0, 2L }), - new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Object[] { 35.0, 18.0, 3L }) + new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Object[] { 21.0, 14.0, 2L, 8.0, 20.0 }), + new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Object[] { 35.0, 18.0, 3L, 6.0, 24.0 }) ); Iterator expectedStarTreeDocumentIterator = inorderStarTreeDocuments.iterator(); StarTreeDocument[] segmentStarTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; for (int i = 0; i < noOfStarTreeDocuments; i++) { - Long metric1 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[0]); + long metric1 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[0]); Long metric2 = starTreeDocuments[i].metrics[1] != null ? NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[1]) : null; - Long metric3 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[2]); - segmentStarTreeDocuments[i] = new StarTreeDocument(starTreeDocuments[i].dimensions, new Object[] { metric1, metric2, metric3 }); + long metric3 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[2]); + long metric4 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[3]); + long metric5 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[4]); + segmentStarTreeDocuments[i] = new StarTreeDocument( + starTreeDocuments[i].dimensions, + new Object[] { metric1, metric2, metric3, metric4, metric5 } + ); } SequentialDocValuesIterator[] dimsIterators = getDimensionIterators(segmentStarTreeDocuments); List metricsIterators = getMetricIterators(segmentStarTreeDocuments); @@ -306,6 +324,8 @@ public void test_sortAndAggregateStarTreeDocuments_nullMetric() throws IOExcepti assertEquals(expectedStarTreeDocument.metrics[0], resultStarTreeDocument.metrics[0]); assertEquals(expectedStarTreeDocument.metrics[1], resultStarTreeDocument.metrics[1]); assertEquals(expectedStarTreeDocument.metrics[2], resultStarTreeDocument.metrics[2]); + assertEquals(expectedStarTreeDocument.metrics[3], resultStarTreeDocument.metrics[3]); + assertEquals(expectedStarTreeDocument.metrics[4], resultStarTreeDocument.metrics[4]); } } @@ -314,26 +334,31 @@ public void test_sortAndAggregateStarTreeDocuments_nullMetricField() throws IOEx int noOfStarTreeDocuments = 5; StarTreeDocument[] starTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; // Setting second metric iterator as empty sorted numeric , indicating a metric field is null - starTreeDocuments[0] = new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Double[] { 12.0, null, randomDouble() }); - starTreeDocuments[1] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Double[] { 10.0, null, randomDouble() }); - starTreeDocuments[2] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Double[] { 14.0, null, randomDouble() }); - starTreeDocuments[3] = new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Double[] { 9.0, null, randomDouble() }); - starTreeDocuments[4] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Double[] { 11.0, null, randomDouble() }); + starTreeDocuments[0] = new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Double[] { 12.0, null, randomDouble(), 8.0, 20.0 }); + starTreeDocuments[1] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Double[] { 10.0, null, randomDouble(), 12.0, 10.0 }); + starTreeDocuments[2] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Double[] { 14.0, null, randomDouble(), 6.0, 24.0 }); + starTreeDocuments[3] = new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Double[] { 9.0, null, randomDouble(), 9.0, 12.0 }); + starTreeDocuments[4] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Double[] { 11.0, null, randomDouble(), 8.0, 13.0 }); List inorderStarTreeDocuments = List.of( - new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Object[] { 21.0, 0.0, 2L }), - new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Object[] { 35.0, 0.0, 3L }) + new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Object[] { 21.0, 0.0, 2L, 8.0, 20.0 }), + new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Object[] { 35.0, 0.0, 3L, 6.0, 24.0 }) ); Iterator expectedStarTreeDocumentIterator = inorderStarTreeDocuments.iterator(); StarTreeDocument[] segmentStarTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; for (int i = 0; i < noOfStarTreeDocuments; i++) { - Long metric1 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[0]); + long metric1 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[0]); Long metric2 = starTreeDocuments[i].metrics[1] != null ? NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[1]) : null; - Long metric3 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[2]); - segmentStarTreeDocuments[i] = new StarTreeDocument(starTreeDocuments[i].dimensions, new Object[] { metric1, metric2, metric3 }); + long metric3 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[2]); + long metric4 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[3]); + long metric5 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[4]); + segmentStarTreeDocuments[i] = new StarTreeDocument( + starTreeDocuments[i].dimensions, + new Object[] { metric1, metric2, metric3, metric4, metric5 } + ); } SequentialDocValuesIterator[] dimsIterators = getDimensionIterators(segmentStarTreeDocuments); List metricsIterators = getMetricIterators(segmentStarTreeDocuments); @@ -353,6 +378,8 @@ public void test_sortAndAggregateStarTreeDocuments_nullMetricField() throws IOEx assertEquals(expectedStarTreeDocument.metrics[0], resultStarTreeDocument.metrics[0]); assertEquals(expectedStarTreeDocument.metrics[1], resultStarTreeDocument.metrics[1]); assertEquals(expectedStarTreeDocument.metrics[2], resultStarTreeDocument.metrics[2]); + assertEquals(expectedStarTreeDocument.metrics[3], resultStarTreeDocument.metrics[3]); + assertEquals(expectedStarTreeDocument.metrics[4], resultStarTreeDocument.metrics[4]); } } @@ -361,11 +388,20 @@ public void test_sortAndAggregateStarTreeDocuments_nullAndMinusOneInDimensionFie int noOfStarTreeDocuments = 5; StarTreeDocument[] starTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; // Setting second metric iterator as empty sorted numeric , indicating a metric field is null - starTreeDocuments[0] = new StarTreeDocument(new Long[] { 2L, null, 3L, 4L }, new Double[] { 12.0, null, randomDouble() }); - starTreeDocuments[1] = new StarTreeDocument(new Long[] { null, 4L, 2L, 1L }, new Double[] { 10.0, null, randomDouble() }); - starTreeDocuments[2] = new StarTreeDocument(new Long[] { null, 4L, 2L, 1L }, new Double[] { 14.0, null, randomDouble() }); - starTreeDocuments[3] = new StarTreeDocument(new Long[] { 2L, null, 3L, 4L }, new Double[] { 9.0, null, randomDouble() }); - starTreeDocuments[4] = new StarTreeDocument(new Long[] { -1L, 4L, 2L, 1L }, new Double[] { 11.0, null, randomDouble() }); + starTreeDocuments[0] = new StarTreeDocument( + new Long[] { 2L, null, 3L, 4L }, + new Double[] { 12.0, null, randomDouble(), 8.0, 20.0 } + ); + starTreeDocuments[1] = new StarTreeDocument( + new Long[] { null, 4L, 2L, 1L }, + new Double[] { 10.0, null, randomDouble(), 12.0, 10.0 } + ); + starTreeDocuments[2] = new StarTreeDocument( + new Long[] { null, 4L, 2L, 1L }, + new Double[] { 14.0, null, randomDouble(), 6.0, 24.0 } + ); + starTreeDocuments[3] = new StarTreeDocument(new Long[] { 2L, null, 3L, 4L }, new Double[] { 9.0, null, randomDouble(), 9.0, 12.0 }); + starTreeDocuments[4] = new StarTreeDocument(new Long[] { -1L, 4L, 2L, 1L }, new Double[] { 11.0, null, randomDouble(), 8.0, 13.0 }); List inorderStarTreeDocuments = List.of( new StarTreeDocument(new Long[] { 2L, null, 3L, 4L }, new Object[] { 21.0, 0.0, 2L }), @@ -376,12 +412,17 @@ public void test_sortAndAggregateStarTreeDocuments_nullAndMinusOneInDimensionFie StarTreeDocument[] segmentStarTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; for (int i = 0; i < noOfStarTreeDocuments; i++) { - Long metric1 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[0]); + long metric1 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[0]); Long metric2 = starTreeDocuments[i].metrics[1] != null ? NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[1]) : null; - Long metric3 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[2]); - segmentStarTreeDocuments[i] = new StarTreeDocument(starTreeDocuments[i].dimensions, new Object[] { metric1, metric2, metric3 }); + long metric3 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[2]); + long metric4 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[3]); + long metric5 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[4]); + segmentStarTreeDocuments[i] = new StarTreeDocument( + starTreeDocuments[i].dimensions, + new Object[] { metric1, metric2, metric3, metric4, metric5 } + ); } SequentialDocValuesIterator[] dimsIterators = getDimensionIterators(segmentStarTreeDocuments); List metricsIterators = getMetricIterators(segmentStarTreeDocuments); @@ -400,6 +441,8 @@ public void test_sortAndAggregateStarTreeDocuments_nullAndMinusOneInDimensionFie assertEquals(expectedStarTreeDocument.metrics[0], resultStarTreeDocument.metrics[0]); assertEquals(expectedStarTreeDocument.metrics[1], resultStarTreeDocument.metrics[1]); assertEquals(expectedStarTreeDocument.metrics[2], resultStarTreeDocument.metrics[2]); + assertEquals(expectedStarTreeDocument.metrics[3], resultStarTreeDocument.metrics[3]); + assertEquals(expectedStarTreeDocument.metrics[4], resultStarTreeDocument.metrics[4]); } builder.build(segmentStarTreeDocumentIterator); validateStarTree(builder.getRootNode(), 4, 1, builder.getStarTreeDocuments()); @@ -409,29 +452,114 @@ public void test_sortAndAggregateStarTreeDocuments_nullDimensionsAndNullMetrics( int noOfStarTreeDocuments = 5; StarTreeDocument[] starTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; // Setting second metric iterator as empty sorted numeric , indicating a metric field is null - starTreeDocuments[0] = new StarTreeDocument(new Long[] { null, null, null, null }, new Double[] { null, null, null }); - starTreeDocuments[1] = new StarTreeDocument(new Long[] { null, null, null, null }, new Double[] { null, null, null }); - starTreeDocuments[2] = new StarTreeDocument(new Long[] { null, null, null, null }, new Double[] { null, null, null }); - starTreeDocuments[3] = new StarTreeDocument(new Long[] { null, null, null, null }, new Double[] { null, null, null }); - starTreeDocuments[4] = new StarTreeDocument(new Long[] { null, null, null, null }, new Double[] { null, null, null }); + starTreeDocuments[0] = new StarTreeDocument(new Long[] { null, null, null, null }, new Double[] { null, null, null, null, null }); + starTreeDocuments[1] = new StarTreeDocument(new Long[] { null, null, null, null }, new Double[] { null, null, null, null, null }); + starTreeDocuments[2] = new StarTreeDocument(new Long[] { null, null, null, null }, new Double[] { null, null, null, null, null }); + starTreeDocuments[3] = new StarTreeDocument(new Long[] { null, null, null, null }, new Double[] { null, null, null, null, null }); + starTreeDocuments[4] = new StarTreeDocument(new Long[] { null, null, null, null }, new Double[] { null, null, null, null, null }); + + List inorderStarTreeDocuments = List.of( + new StarTreeDocument(new Long[] { null, null, null, null }, new Object[] { 0.0, 0.0, 0L, null, null }) + ); + Iterator expectedStarTreeDocumentIterator = inorderStarTreeDocuments.iterator(); + + StarTreeDocument[] segmentStarTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; + for (int i = 0; i < noOfStarTreeDocuments; i++) { + Long metric1 = starTreeDocuments[i].metrics[0] != null + ? NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[0]) + : null; + Long metric2 = starTreeDocuments[i].metrics[1] != null + ? NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[1]) + : null; + Long metric3 = starTreeDocuments[i].metrics[2] != null + ? NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[2]) + : null; + Long metric4 = starTreeDocuments[i].metrics[3] != null + ? NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[3]) + : null; + Long metric5 = starTreeDocuments[i].metrics[4] != null + ? NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[4]) + : null; + segmentStarTreeDocuments[i] = new StarTreeDocument( + starTreeDocuments[i].dimensions, + new Object[] { metric1, metric2, metric3, metric4, metric5 } + ); + } + SequentialDocValuesIterator[] dimsIterators = getDimensionIterators(segmentStarTreeDocuments); + List metricsIterators = getMetricIterators(segmentStarTreeDocuments); + builder = getStarTreeBuilder(compositeField, writeState, mapperService); + Iterator segmentStarTreeDocumentIterator = builder.sortAndAggregateSegmentDocuments( + dimsIterators, + metricsIterators + ); + + while (segmentStarTreeDocumentIterator.hasNext() && expectedStarTreeDocumentIterator.hasNext()) { + StarTreeDocument resultStarTreeDocument = segmentStarTreeDocumentIterator.next(); + StarTreeDocument expectedStarTreeDocument = expectedStarTreeDocumentIterator.next(); + assertEquals(expectedStarTreeDocument.dimensions[0], resultStarTreeDocument.dimensions[0]); + assertEquals(expectedStarTreeDocument.dimensions[1], resultStarTreeDocument.dimensions[1]); + assertEquals(expectedStarTreeDocument.dimensions[2], resultStarTreeDocument.dimensions[2]); + assertEquals(expectedStarTreeDocument.dimensions[3], resultStarTreeDocument.dimensions[3]); + assertEquals(expectedStarTreeDocument.metrics[0], resultStarTreeDocument.metrics[0]); + assertEquals(expectedStarTreeDocument.metrics[1], resultStarTreeDocument.metrics[1]); + assertEquals(expectedStarTreeDocument.metrics[2], resultStarTreeDocument.metrics[2]); + assertEquals(expectedStarTreeDocument.metrics[3], resultStarTreeDocument.metrics[3]); + assertEquals(expectedStarTreeDocument.metrics[4], resultStarTreeDocument.metrics[4]); + } + builder.build(segmentStarTreeDocumentIterator); + validateStarTree(builder.getRootNode(), 4, 1, builder.getStarTreeDocuments()); + } + + public void test_sortAndAggregateStarTreeDocuments_nullDimensionsAndFewNullMetrics() throws IOException { + int noOfStarTreeDocuments = 5; + StarTreeDocument[] starTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; + + double sumValue = randomDouble(); + double minValue = randomDouble(); + double maxValue = randomDouble(); + + // Setting second metric iterator as empty sorted numeric , indicating a metric field is null + starTreeDocuments[0] = new StarTreeDocument( + new Long[] { null, null, null, null }, + new Double[] { null, null, randomDouble(), null, maxValue } + ); + starTreeDocuments[1] = new StarTreeDocument(new Long[] { null, null, null, null }, new Double[] { null, null, null, null, null }); + starTreeDocuments[2] = new StarTreeDocument( + new Long[] { null, null, null, null }, + new Double[] { null, null, null, minValue, null } + ); + starTreeDocuments[3] = new StarTreeDocument(new Long[] { null, null, null, null }, new Double[] { null, null, null, null, null }); + starTreeDocuments[4] = new StarTreeDocument( + new Long[] { null, null, null, null }, + new Double[] { sumValue, null, randomDouble(), null, null } + ); List inorderStarTreeDocuments = List.of( - new StarTreeDocument(new Long[] { null, null, null, null }, new Object[] { 0.0, 0.0, 5L }) + new StarTreeDocument(new Long[] { null, null, null, null }, new Object[] { sumValue, 0.0, 2L, minValue, maxValue }) ); Iterator expectedStarTreeDocumentIterator = inorderStarTreeDocuments.iterator(); StarTreeDocument[] segmentStarTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; for (int i = 0; i < noOfStarTreeDocuments; i++) { - Long metric1 = starTreeDocuments[i].metrics[1] != null + Long metric1 = starTreeDocuments[i].metrics[0] != null ? NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[0]) : null; Long metric2 = starTreeDocuments[i].metrics[1] != null ? NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[1]) : null; - Long metric3 = starTreeDocuments[i].metrics[1] != null + Long metric3 = starTreeDocuments[i].metrics[2] != null ? NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[2]) : null; - segmentStarTreeDocuments[i] = new StarTreeDocument(starTreeDocuments[i].dimensions, new Object[] { metric1, metric2, metric3 }); + Long metric4 = starTreeDocuments[i].metrics[3] != null + ? NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[3]) + : null; + Long metric5 = starTreeDocuments[i].metrics[4] != null + ? NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[4]) + : null; + segmentStarTreeDocuments[i] = new StarTreeDocument( + starTreeDocuments[i].dimensions, + new Object[] { metric1, metric2, metric3, metric4, metric5 } + ); } SequentialDocValuesIterator[] dimsIterators = getDimensionIterators(segmentStarTreeDocuments); List metricsIterators = getMetricIterators(segmentStarTreeDocuments); @@ -451,6 +579,8 @@ public void test_sortAndAggregateStarTreeDocuments_nullDimensionsAndNullMetrics( assertEquals(expectedStarTreeDocument.metrics[0], resultStarTreeDocument.metrics[0]); assertEquals(expectedStarTreeDocument.metrics[1], resultStarTreeDocument.metrics[1]); assertEquals(expectedStarTreeDocument.metrics[2], resultStarTreeDocument.metrics[2]); + assertEquals(expectedStarTreeDocument.metrics[3], resultStarTreeDocument.metrics[3]); + assertEquals(expectedStarTreeDocument.metrics[4], resultStarTreeDocument.metrics[4]); } builder.build(segmentStarTreeDocumentIterator); validateStarTree(builder.getRootNode(), 4, 1, builder.getStarTreeDocuments()); @@ -461,14 +591,29 @@ public void test_sortAndAggregateStarTreeDocuments_emptyDimensions() throws IOEx int noOfStarTreeDocuments = 5; StarTreeDocument[] starTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; // Setting second metric iterator as empty sorted numeric , indicating a metric field is null - starTreeDocuments[0] = new StarTreeDocument(new Long[] { null, null, null, null }, new Double[] { 12.0, null, randomDouble() }); - starTreeDocuments[1] = new StarTreeDocument(new Long[] { null, null, null, null }, new Double[] { 10.0, null, randomDouble() }); - starTreeDocuments[2] = new StarTreeDocument(new Long[] { null, null, null, null }, new Double[] { 14.0, null, randomDouble() }); - starTreeDocuments[3] = new StarTreeDocument(new Long[] { null, null, null, null }, new Double[] { 9.0, null, randomDouble() }); - starTreeDocuments[4] = new StarTreeDocument(new Long[] { null, null, null, null }, new Double[] { 11.0, null, randomDouble() }); + starTreeDocuments[0] = new StarTreeDocument( + new Long[] { null, null, null, null }, + new Double[] { 12.0, null, randomDouble(), 8.0, 20.0 } + ); + starTreeDocuments[1] = new StarTreeDocument( + new Long[] { null, null, null, null }, + new Double[] { 10.0, null, randomDouble(), 12.0, 10.0 } + ); + starTreeDocuments[2] = new StarTreeDocument( + new Long[] { null, null, null, null }, + new Double[] { 14.0, null, randomDouble(), 6.0, 24.0 } + ); + starTreeDocuments[3] = new StarTreeDocument( + new Long[] { null, null, null, null }, + new Double[] { 9.0, null, randomDouble(), 9.0, 12.0 } + ); + starTreeDocuments[4] = new StarTreeDocument( + new Long[] { null, null, null, null }, + new Double[] { 11.0, null, randomDouble(), 8.0, 13.0 } + ); List inorderStarTreeDocuments = List.of( - new StarTreeDocument(new Long[] { null, null, null, null }, new Object[] { 56.0, 0.0, 5L }) + new StarTreeDocument(new Long[] { null, null, null, null }, new Object[] { 56.0, 0.0, 5L, 6.0, 24.0 }) ); Iterator expectedStarTreeDocumentIterator = inorderStarTreeDocuments.iterator(); @@ -479,7 +624,12 @@ public void test_sortAndAggregateStarTreeDocuments_emptyDimensions() throws IOEx ? NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[1]) : null; Long metric3 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[2]); - segmentStarTreeDocuments[i] = new StarTreeDocument(starTreeDocuments[i].dimensions, new Object[] { metric1, metric2, metric3 }); + Long metric4 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[3]); + Long metric5 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[4]); + segmentStarTreeDocuments[i] = new StarTreeDocument( + starTreeDocuments[i].dimensions, + new Object[] { metric1, metric2, metric3, metric4, metric5 } + ); } SequentialDocValuesIterator[] dimsIterators = getDimensionIterators(segmentStarTreeDocuments); List metricsIterators = getMetricIterators(segmentStarTreeDocuments); @@ -499,6 +649,8 @@ public void test_sortAndAggregateStarTreeDocuments_emptyDimensions() throws IOEx assertEquals(expectedStarTreeDocument.metrics[0], resultStarTreeDocument.metrics[0]); assertEquals(expectedStarTreeDocument.metrics[1], resultStarTreeDocument.metrics[1]); assertEquals(expectedStarTreeDocument.metrics[2], resultStarTreeDocument.metrics[2]); + assertEquals(expectedStarTreeDocument.metrics[3], resultStarTreeDocument.metrics[3]); + assertEquals(expectedStarTreeDocument.metrics[4], resultStarTreeDocument.metrics[4]); } } @@ -507,15 +659,30 @@ public void test_sortAndAggregateStarTreeDocument_longMaxAndLongMinDimensions() int noOfStarTreeDocuments = 5; StarTreeDocument[] starTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; - starTreeDocuments[0] = new StarTreeDocument(new Long[] { Long.MIN_VALUE, 4L, 3L, 4L }, new Double[] { 12.0, 10.0, randomDouble() }); - starTreeDocuments[1] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, Long.MAX_VALUE }, new Double[] { 10.0, 6.0, randomDouble() }); - starTreeDocuments[2] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, Long.MAX_VALUE }, new Double[] { 14.0, 12.0, randomDouble() }); - starTreeDocuments[3] = new StarTreeDocument(new Long[] { Long.MIN_VALUE, 4L, 3L, 4L }, new Double[] { 9.0, 4.0, randomDouble() }); - starTreeDocuments[4] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, Long.MAX_VALUE }, new Double[] { 11.0, 16.0, randomDouble() }); + starTreeDocuments[0] = new StarTreeDocument( + new Long[] { Long.MIN_VALUE, 4L, 3L, 4L }, + new Double[] { 12.0, 10.0, randomDouble(), 8.0, 20.0 } + ); + starTreeDocuments[1] = new StarTreeDocument( + new Long[] { 3L, 4L, 2L, Long.MAX_VALUE }, + new Double[] { 10.0, 6.0, randomDouble(), 12.0, 10.0 } + ); + starTreeDocuments[2] = new StarTreeDocument( + new Long[] { 3L, 4L, 2L, Long.MAX_VALUE }, + new Double[] { 14.0, 12.0, randomDouble(), 6.0, 24.0 } + ); + starTreeDocuments[3] = new StarTreeDocument( + new Long[] { Long.MIN_VALUE, 4L, 3L, 4L }, + new Double[] { 9.0, 4.0, randomDouble(), 9.0, 12.0 } + ); + starTreeDocuments[4] = new StarTreeDocument( + new Long[] { 3L, 4L, 2L, Long.MAX_VALUE }, + new Double[] { 11.0, 16.0, randomDouble(), 8.0, 13.0 } + ); List inorderStarTreeDocuments = List.of( - new StarTreeDocument(new Long[] { Long.MIN_VALUE, 4L, 3L, 4L }, new Object[] { 21.0, 14.0, 2L }), - new StarTreeDocument(new Long[] { 3L, 4L, 2L, Long.MAX_VALUE }, new Object[] { 35.0, 34.0, 3L }) + new StarTreeDocument(new Long[] { Long.MIN_VALUE, 4L, 3L, 4L }, new Object[] { 21.0, 14.0, 2L, 8.0, 20.0 }), + new StarTreeDocument(new Long[] { 3L, 4L, 2L, Long.MAX_VALUE }, new Object[] { 35.0, 34.0, 3L, 6.0, 24.0 }) ); Iterator expectedStarTreeDocumentIterator = inorderStarTreeDocuments.iterator(); @@ -524,7 +691,12 @@ public void test_sortAndAggregateStarTreeDocument_longMaxAndLongMinDimensions() long metric1 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[0]); long metric2 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[1]); long metric3 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[2]); - segmentStarTreeDocuments[i] = new StarTreeDocument(starTreeDocuments[i].dimensions, new Long[] { metric1, metric2, metric3 }); + long metric4 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[3]); + long metric5 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[4]); + segmentStarTreeDocuments[i] = new StarTreeDocument( + starTreeDocuments[i].dimensions, + new Long[] { metric1, metric2, metric3, metric4, metric5 } + ); } SequentialDocValuesIterator[] dimsIterators = getDimensionIterators(segmentStarTreeDocuments); @@ -546,6 +718,8 @@ public void test_sortAndAggregateStarTreeDocument_longMaxAndLongMinDimensions() assertEquals(expectedStarTreeDocument.metrics[0], resultStarTreeDocument.metrics[0]); assertEquals(expectedStarTreeDocument.metrics[1], resultStarTreeDocument.metrics[1]); assertEquals(expectedStarTreeDocument.metrics[2], resultStarTreeDocument.metrics[2]); + assertEquals(expectedStarTreeDocument.metrics[3], resultStarTreeDocument.metrics[3]); + assertEquals(expectedStarTreeDocument.metrics[4], resultStarTreeDocument.metrics[4]); numOfAggregatedDocuments++; } @@ -559,15 +733,21 @@ public void test_sortAndAggregateStarTreeDocument_DoubleMaxAndDoubleMinMetrics() int noOfStarTreeDocuments = 5; StarTreeDocument[] starTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; - starTreeDocuments[0] = new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Double[] { Double.MAX_VALUE, 10.0, randomDouble() }); - starTreeDocuments[1] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Double[] { 10.0, 6.0, randomDouble() }); - starTreeDocuments[2] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Double[] { 14.0, Double.MIN_VALUE, randomDouble() }); - starTreeDocuments[3] = new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Double[] { 9.0, 4.0, randomDouble() }); - starTreeDocuments[4] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Double[] { 11.0, 16.0, randomDouble() }); + starTreeDocuments[0] = new StarTreeDocument( + new Long[] { 2L, 4L, 3L, 4L }, + new Double[] { Double.MAX_VALUE, 10.0, randomDouble(), 8.0, 20.0 } + ); + starTreeDocuments[1] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Double[] { 10.0, 6.0, randomDouble(), 12.0, 10.0 }); + starTreeDocuments[2] = new StarTreeDocument( + new Long[] { 3L, 4L, 2L, 1L }, + new Double[] { 14.0, Double.MIN_VALUE, randomDouble(), 6.0, 24.0 } + ); + starTreeDocuments[3] = new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Double[] { 9.0, 4.0, randomDouble(), 9.0, 12.0 }); + starTreeDocuments[4] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Double[] { 11.0, 16.0, randomDouble(), 8.0, 13.0 }); List inorderStarTreeDocuments = List.of( - new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Object[] { Double.MAX_VALUE + 9, 14.0, 2L }), - new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Object[] { 35.0, Double.MIN_VALUE + 22, 3L }) + new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Object[] { Double.MAX_VALUE + 9, 14.0, 2L, 8.0, 20.0 }), + new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Object[] { 35.0, Double.MIN_VALUE + 22, 3L, 6.0, 24.0 }) ); Iterator expectedStarTreeDocumentIterator = inorderStarTreeDocuments.iterator(); @@ -576,7 +756,12 @@ public void test_sortAndAggregateStarTreeDocument_DoubleMaxAndDoubleMinMetrics() long metric1 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[0]); long metric2 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[1]); long metric3 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[2]); - segmentStarTreeDocuments[i] = new StarTreeDocument(starTreeDocuments[i].dimensions, new Long[] { metric1, metric2, metric3 }); + long metric4 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[3]); + long metric5 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[4]); + segmentStarTreeDocuments[i] = new StarTreeDocument( + starTreeDocuments[i].dimensions, + new Long[] { metric1, metric2, metric3, metric4, metric5 } + ); } SequentialDocValuesIterator[] dimsIterators = getDimensionIterators(segmentStarTreeDocuments); @@ -598,6 +783,8 @@ public void test_sortAndAggregateStarTreeDocument_DoubleMaxAndDoubleMinMetrics() assertEquals(expectedStarTreeDocument.metrics[0], resultStarTreeDocument.metrics[0]); assertEquals(expectedStarTreeDocument.metrics[1], resultStarTreeDocument.metrics[1]); assertEquals(expectedStarTreeDocument.metrics[2], resultStarTreeDocument.metrics[2]); + assertEquals(expectedStarTreeDocument.metrics[3], resultStarTreeDocument.metrics[3]); + assertEquals(expectedStarTreeDocument.metrics[4], resultStarTreeDocument.metrics[4]); numOfAggregatedDocuments++; } @@ -620,8 +807,16 @@ public void test_build_halfFloatMetrics() throws IOException { .build(new Mapper.BuilderContext(settings, new ContentPath())); NumberFieldMapper numberFieldMapper3 = new NumberFieldMapper.Builder("field6", NumberFieldMapper.NumberType.HALF_FLOAT, false, true) .build(new Mapper.BuilderContext(settings, new ContentPath())); + NumberFieldMapper numberFieldMapper4 = new NumberFieldMapper.Builder("field9", NumberFieldMapper.NumberType.HALF_FLOAT, false, true) + .build(new Mapper.BuilderContext(settings, new ContentPath())); + NumberFieldMapper numberFieldMapper5 = new NumberFieldMapper.Builder( + "field10", + NumberFieldMapper.NumberType.HALF_FLOAT, + false, + true + ).build(new Mapper.BuilderContext(settings, new ContentPath())); MappingLookup fieldMappers = new MappingLookup( - Set.of(numberFieldMapper1, numberFieldMapper2, numberFieldMapper3), + Set.of(numberFieldMapper1, numberFieldMapper2, numberFieldMapper3, numberFieldMapper4, numberFieldMapper5), Collections.emptyList(), Collections.emptyList(), 0, @@ -634,23 +829,48 @@ public void test_build_halfFloatMetrics() throws IOException { starTreeDocuments[0] = new StarTreeDocument( new Long[] { 2L, 4L, 3L, 4L }, - new HalfFloatPoint[] { new HalfFloatPoint("hf1", 12), new HalfFloatPoint("hf6", 10), new HalfFloatPoint("field6", 10) } + new HalfFloatPoint[] { + new HalfFloatPoint("hf1", 12), + new HalfFloatPoint("hf6", 10), + new HalfFloatPoint("field6", 10), + new HalfFloatPoint("field9", 8), + new HalfFloatPoint("field10", 20) } ); starTreeDocuments[1] = new StarTreeDocument( new Long[] { 3L, 4L, 2L, 1L }, - new HalfFloatPoint[] { new HalfFloatPoint("hf2", 10), new HalfFloatPoint("hf7", 6), new HalfFloatPoint("field6", 10) } + new HalfFloatPoint[] { + new HalfFloatPoint("hf2", 10), + new HalfFloatPoint("hf7", 6), + new HalfFloatPoint("field6", 10), + new HalfFloatPoint("field9", 12), + new HalfFloatPoint("field10", 10) } ); starTreeDocuments[2] = new StarTreeDocument( new Long[] { 3L, 4L, 2L, 1L }, - new HalfFloatPoint[] { new HalfFloatPoint("hf3", 14), new HalfFloatPoint("hf8", 12), new HalfFloatPoint("field6", 10) } + new HalfFloatPoint[] { + new HalfFloatPoint("hf3", 14), + new HalfFloatPoint("hf8", 12), + new HalfFloatPoint("field6", 10), + new HalfFloatPoint("field9", 6), + new HalfFloatPoint("field10", 24) } ); starTreeDocuments[3] = new StarTreeDocument( new Long[] { 2L, 4L, 3L, 4L }, - new HalfFloatPoint[] { new HalfFloatPoint("hf4", 9), new HalfFloatPoint("hf9", 4), new HalfFloatPoint("field6", 10) } + new HalfFloatPoint[] { + new HalfFloatPoint("hf4", 9), + new HalfFloatPoint("hf9", 4), + new HalfFloatPoint("field6", 10), + new HalfFloatPoint("field9", 9), + new HalfFloatPoint("field10", 12) } ); starTreeDocuments[4] = new StarTreeDocument( new Long[] { 3L, 4L, 2L, 1L }, - new HalfFloatPoint[] { new HalfFloatPoint("hf5", 11), new HalfFloatPoint("hf10", 16), new HalfFloatPoint("field6", 10) } + new HalfFloatPoint[] { + new HalfFloatPoint("hf5", 11), + new HalfFloatPoint("hf10", 16), + new HalfFloatPoint("field6", 10), + new HalfFloatPoint("field9", 8), + new HalfFloatPoint("field10", 13) } ); StarTreeDocument[] segmentStarTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; @@ -664,7 +884,16 @@ public void test_build_halfFloatMetrics() throws IOException { long metric3 = HalfFloatPoint.halfFloatToSortableShort( ((HalfFloatPoint) starTreeDocuments[i].metrics[2]).numericValue().floatValue() ); - segmentStarTreeDocuments[i] = new StarTreeDocument(starTreeDocuments[i].dimensions, new Long[] { metric1, metric2, metric3 }); + long metric4 = HalfFloatPoint.halfFloatToSortableShort( + ((HalfFloatPoint) starTreeDocuments[i].metrics[3]).numericValue().floatValue() + ); + long metric5 = HalfFloatPoint.halfFloatToSortableShort( + ((HalfFloatPoint) starTreeDocuments[i].metrics[4]).numericValue().floatValue() + ); + segmentStarTreeDocuments[i] = new StarTreeDocument( + starTreeDocuments[i].dimensions, + new Long[] { metric1, metric2, metric3, metric4, metric5 } + ); } SequentialDocValuesIterator[] dimsIterators = getDimensionIterators(segmentStarTreeDocuments); @@ -696,8 +925,12 @@ public void test_build_floatMetrics() throws IOException { .build(new Mapper.BuilderContext(settings, new ContentPath())); NumberFieldMapper numberFieldMapper3 = new NumberFieldMapper.Builder("field6", NumberFieldMapper.NumberType.FLOAT, false, true) .build(new Mapper.BuilderContext(settings, new ContentPath())); + NumberFieldMapper numberFieldMapper4 = new NumberFieldMapper.Builder("field9", NumberFieldMapper.NumberType.FLOAT, false, true) + .build(new Mapper.BuilderContext(settings, new ContentPath())); + NumberFieldMapper numberFieldMapper5 = new NumberFieldMapper.Builder("field10", NumberFieldMapper.NumberType.FLOAT, false, true) + .build(new Mapper.BuilderContext(settings, new ContentPath())); MappingLookup fieldMappers = new MappingLookup( - Set.of(numberFieldMapper1, numberFieldMapper2, numberFieldMapper3), + Set.of(numberFieldMapper1, numberFieldMapper2, numberFieldMapper3, numberFieldMapper4, numberFieldMapper5), Collections.emptyList(), Collections.emptyList(), 0, @@ -708,18 +941,35 @@ public void test_build_floatMetrics() throws IOException { int noOfStarTreeDocuments = 5; StarTreeDocument[] starTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; - starTreeDocuments[0] = new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Float[] { 12.0F, 10.0F, randomFloat() }); - starTreeDocuments[1] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Float[] { 10.0F, 6.0F, randomFloat() }); - starTreeDocuments[2] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Float[] { 14.0F, 12.0F, randomFloat() }); - starTreeDocuments[3] = new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Float[] { 9.0F, 4.0F, randomFloat() }); - starTreeDocuments[4] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Float[] { 11.0F, 16.0F, randomFloat() }); + starTreeDocuments[0] = new StarTreeDocument( + new Long[] { 2L, 4L, 3L, 4L }, + new Float[] { 12.0F, 10.0F, randomFloat(), 8.0F, 20.0F } + ); + starTreeDocuments[1] = new StarTreeDocument( + new Long[] { 3L, 4L, 2L, 1L }, + new Float[] { 10.0F, 6.0F, randomFloat(), 12.0F, 10.0F } + ); + starTreeDocuments[2] = new StarTreeDocument( + new Long[] { 3L, 4L, 2L, 1L }, + new Float[] { 14.0F, 12.0F, randomFloat(), 6.0F, 24.0F } + ); + starTreeDocuments[3] = new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Float[] { 9.0F, 4.0F, randomFloat(), 9.0F, 12.0F }); + starTreeDocuments[4] = new StarTreeDocument( + new Long[] { 3L, 4L, 2L, 1L }, + new Float[] { 11.0F, 16.0F, randomFloat(), 8.0F, 13.0F } + ); StarTreeDocument[] segmentStarTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; for (int i = 0; i < noOfStarTreeDocuments; i++) { long metric1 = NumericUtils.floatToSortableInt((Float) starTreeDocuments[i].metrics[0]); long metric2 = NumericUtils.floatToSortableInt((Float) starTreeDocuments[i].metrics[1]); long metric3 = NumericUtils.floatToSortableInt((Float) starTreeDocuments[i].metrics[2]); - segmentStarTreeDocuments[i] = new StarTreeDocument(starTreeDocuments[i].dimensions, new Long[] { metric1, metric2, metric3 }); + long metric4 = NumericUtils.floatToSortableInt((Float) starTreeDocuments[i].metrics[3]); + long metric5 = NumericUtils.floatToSortableInt((Float) starTreeDocuments[i].metrics[4]); + segmentStarTreeDocuments[i] = new StarTreeDocument( + starTreeDocuments[i].dimensions, + new Long[] { metric1, metric2, metric3, metric4, metric5 } + ); } SequentialDocValuesIterator[] dimsIterators = getDimensionIterators(segmentStarTreeDocuments); @@ -750,8 +1000,12 @@ public void test_build_longMetrics() throws IOException { .build(new Mapper.BuilderContext(settings, new ContentPath())); NumberFieldMapper numberFieldMapper3 = new NumberFieldMapper.Builder("field6", NumberFieldMapper.NumberType.LONG, false, true) .build(new Mapper.BuilderContext(settings, new ContentPath())); + NumberFieldMapper numberFieldMapper4 = new NumberFieldMapper.Builder("field9", NumberFieldMapper.NumberType.LONG, false, true) + .build(new Mapper.BuilderContext(settings, new ContentPath())); + NumberFieldMapper numberFieldMapper5 = new NumberFieldMapper.Builder("field10", NumberFieldMapper.NumberType.LONG, false, true) + .build(new Mapper.BuilderContext(settings, new ContentPath())); MappingLookup fieldMappers = new MappingLookup( - Set.of(numberFieldMapper1, numberFieldMapper2, numberFieldMapper3), + Set.of(numberFieldMapper1, numberFieldMapper2, numberFieldMapper3, numberFieldMapper4, numberFieldMapper5), Collections.emptyList(), Collections.emptyList(), 0, @@ -762,18 +1016,23 @@ public void test_build_longMetrics() throws IOException { int noOfStarTreeDocuments = 5; StarTreeDocument[] starTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; - starTreeDocuments[0] = new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Long[] { 12L, 10L, randomLong() }); - starTreeDocuments[1] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Long[] { 10L, 6L, randomLong() }); - starTreeDocuments[2] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Long[] { 14L, 12L, randomLong() }); - starTreeDocuments[3] = new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Long[] { 9L, 4L, randomLong() }); - starTreeDocuments[4] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Long[] { 11L, 16L, randomLong() }); + starTreeDocuments[0] = new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Long[] { 12L, 10L, randomLong(), 8L, 20L }); + starTreeDocuments[1] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Long[] { 10L, 6L, randomLong(), 12L, 10L }); + starTreeDocuments[2] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Long[] { 14L, 12L, randomLong(), 6L, 24L }); + starTreeDocuments[3] = new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Long[] { 9L, 4L, randomLong(), 9L, 12L }); + starTreeDocuments[4] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Long[] { 11L, 16L, randomLong(), 8L, 13L }); StarTreeDocument[] segmentStarTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; for (int i = 0; i < noOfStarTreeDocuments; i++) { long metric1 = (Long) starTreeDocuments[i].metrics[0]; long metric2 = (Long) starTreeDocuments[i].metrics[1]; long metric3 = (Long) starTreeDocuments[i].metrics[2]; - segmentStarTreeDocuments[i] = new StarTreeDocument(starTreeDocuments[i].dimensions, new Long[] { metric1, metric2, metric3 }); + long metric4 = (Long) starTreeDocuments[i].metrics[3]; + long metric5 = (Long) starTreeDocuments[i].metrics[4]; + segmentStarTreeDocuments[i] = new StarTreeDocument( + starTreeDocuments[i].dimensions, + new Long[] { metric1, metric2, metric3, metric4, metric5 } + ); } SequentialDocValuesIterator[] dimsIterators = getDimensionIterators(segmentStarTreeDocuments); @@ -794,14 +1053,13 @@ public void test_build_longMetrics() throws IOException { private static Iterator getExpectedStarTreeDocumentIterator() { List expectedStarTreeDocuments = List.of( - new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Object[] { 21.0, 14.0, 2L }), - new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Object[] { 35.0, 34.0, 3L }), - new StarTreeDocument(new Long[] { null, 4L, 2L, 1L }, new Object[] { 35.0, 34.0, 3L }), - new StarTreeDocument(new Long[] { null, 4L, 3L, 4L }, new Object[] { 21.0, 14.0, 2L }), - new StarTreeDocument(new Long[] { null, 4L, null, 1L }, new Object[] { 35.0, 34.0, 3L }), - new StarTreeDocument(new Long[] { null, 4L, null, 4L }, new Object[] { 21.0, 14.0, 2L }), - new StarTreeDocument(new Long[] { null, 4L, null, null }, new Object[] { 56.0, 48.0, 5L }), - new StarTreeDocument(new Long[] { null, null, null, null }, new Object[] { 56.0, 48.0, 5L }) + new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Object[] { 21.0, 14.0, 2L, 8.0, 20.0 }), + new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Object[] { 35.0, 34.0, 3L, 6.0, 24.0 }), + new StarTreeDocument(new Long[] { null, 4L, 2L, 1L }, new Object[] { 35.0, 34.0, 3L, 6.0, 24.0 }), + new StarTreeDocument(new Long[] { null, 4L, 3L, 4L }, new Object[] { 21.0, 14.0, 2L, 8.0, 20.0 }), + new StarTreeDocument(new Long[] { null, 4L, null, 1L }, new Object[] { 35.0, 34.0, 3L, 6.0, 24.0 }), + new StarTreeDocument(new Long[] { null, 4L, null, 4L }, new Object[] { 21.0, 14.0, 2L, 8.0, 20.0 }), + new StarTreeDocument(new Long[] { null, 4L, null, null }, new Object[] { 56.0, 48.0, 5L, 6.0, 24.0 }) ); return expectedStarTreeDocuments.iterator(); } @@ -811,18 +1069,23 @@ public void test_build() throws IOException { int noOfStarTreeDocuments = 5; StarTreeDocument[] starTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; - starTreeDocuments[0] = new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Double[] { 12.0, 10.0, randomDouble() }); - starTreeDocuments[1] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Double[] { 10.0, 6.0, randomDouble() }); - starTreeDocuments[2] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Double[] { 14.0, 12.0, randomDouble() }); - starTreeDocuments[3] = new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Double[] { 9.0, 4.0, randomDouble() }); - starTreeDocuments[4] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Double[] { 11.0, 16.0, randomDouble() }); + starTreeDocuments[0] = new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Double[] { 12.0, 10.0, randomDouble(), 8.0, 20.0 }); + starTreeDocuments[1] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Double[] { 10.0, 6.0, randomDouble(), 12.0, 10.0 }); + starTreeDocuments[2] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Double[] { 14.0, 12.0, randomDouble(), 6.0, 24.0 }); + starTreeDocuments[3] = new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Double[] { 9.0, 4.0, randomDouble(), 9.0, 12.0 }); + starTreeDocuments[4] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Double[] { 11.0, 16.0, randomDouble(), 8.0, 13.0 }); StarTreeDocument[] segmentStarTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; for (int i = 0; i < noOfStarTreeDocuments; i++) { long metric1 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[0]); long metric2 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[1]); long metric3 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[2]); - segmentStarTreeDocuments[i] = new StarTreeDocument(starTreeDocuments[i].dimensions, new Long[] { metric1, metric2, metric3 }); + long metric4 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[3]); + long metric5 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[4]); + segmentStarTreeDocuments[i] = new StarTreeDocument( + starTreeDocuments[i].dimensions, + new Long[] { metric1, metric2, metric3, metric4, metric5 } + ); } SequentialDocValuesIterator[] dimsIterators = getDimensionIterators(segmentStarTreeDocuments); @@ -857,6 +1120,8 @@ private void assertStarTreeDocuments( assertEquals(expectedStarTreeDocument.metrics[0], resultStarTreeDocument.metrics[0]); assertEquals(expectedStarTreeDocument.metrics[1], resultStarTreeDocument.metrics[1]); assertEquals(expectedStarTreeDocument.metrics[2], resultStarTreeDocument.metrics[2]); + assertEquals(expectedStarTreeDocument.metrics[3], resultStarTreeDocument.metrics[3]); + assertEquals(expectedStarTreeDocument.metrics[4], resultStarTreeDocument.metrics[4]); } } From 2d13df7f9ef4aff5136720929b814fe9da63776c Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Mon, 12 Aug 2024 11:08:51 -0400 Subject: [PATCH 19/40] Bump org.xerial.snappy:snappy-java from 1.1.10.5 to 1.1.10.6 in /test/fixtures/hdfs-fixture (#15207) * Bump org.xerial.snappy:snappy-java in /test/fixtures/hdfs-fixture Bumps [org.xerial.snappy:snappy-java](https://github.com/xerial/snappy-java) from 1.1.10.5 to 1.1.10.6. - [Release notes](https://github.com/xerial/snappy-java/releases) - [Commits](https://github.com/xerial/snappy-java/compare/v1.1.10.5...v1.1.10.6) --- updated-dependencies: - dependency-name: org.xerial.snappy:snappy-java dependency-type: direct:production update-type: version-update:semver-patch ... Signed-off-by: dependabot[bot] * Update changelog Signed-off-by: dependabot[bot] --------- Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> Co-authored-by: dependabot[bot] --- CHANGELOG.md | 1 + test/fixtures/hdfs-fixture/build.gradle | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 7fd56693e1ea6..c60b8f492d02c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -29,6 +29,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), - Bump `org.apache.avro:avro` from 1.11.3 to 1.12.0 in /plugins/repository-hdfs ([#15119](https://github.com/opensearch-project/OpenSearch/pull/15119)) - Bump `org.bouncycastle:bcpg-fips` from 1.0.7.1 to 2.0.8 and `org.bouncycastle:bc-fips` from 1.0.2.5 to 2.0.0 in /distribution/tools/plugin-cli ([#15103](https://github.com/opensearch-project/OpenSearch/pull/15103)) - Bump `com.azure:azure-core` from 1.49.1 to 1.51.0 ([#15111](https://github.com/opensearch-project/OpenSearch/pull/15111)) +- Bump `org.xerial.snappy:snappy-java` from 1.1.10.5 to 1.1.10.6 ([#15207](https://github.com/opensearch-project/OpenSearch/pull/15207)) ### Changed - Add lower limit for primary and replica batch allocators timeout ([#14979](https://github.com/opensearch-project/OpenSearch/pull/14979)) diff --git a/test/fixtures/hdfs-fixture/build.gradle b/test/fixtures/hdfs-fixture/build.gradle index 9b8f62b8c55b8..aaa150e73a13e 100644 --- a/test/fixtures/hdfs-fixture/build.gradle +++ b/test/fixtures/hdfs-fixture/build.gradle @@ -88,5 +88,5 @@ dependencies { exclude group: "com.squareup.okio" } runtimeOnly "com.squareup.okio:okio:3.9.0" - runtimeOnly "org.xerial.snappy:snappy-java:1.1.10.5" + runtimeOnly "org.xerial.snappy:snappy-java:1.1.10.6" } From 2b5c6b76458d24c7aee1af1b91c32c8e7a16b665 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Mon, 12 Aug 2024 11:11:11 -0400 Subject: [PATCH 20/40] Bump org.apache.commons:commons-lang3 from 3.15.0 to 3.16.0 in /plugins/repository-hdfs (#15205) * Bump org.apache.commons:commons-lang3 in /plugins/repository-hdfs Bumps org.apache.commons:commons-lang3 from 3.15.0 to 3.16.0. --- updated-dependencies: - dependency-name: org.apache.commons:commons-lang3 dependency-type: direct:production update-type: version-update:semver-minor ... Signed-off-by: dependabot[bot] * Updating SHAs Signed-off-by: dependabot[bot] * Update changelog Signed-off-by: dependabot[bot] --------- Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> Co-authored-by: dependabot[bot] --- CHANGELOG.md | 2 +- plugins/repository-hdfs/build.gradle | 2 +- plugins/repository-hdfs/licenses/commons-lang3-3.15.0.jar.sha1 | 1 - plugins/repository-hdfs/licenses/commons-lang3-3.16.0.jar.sha1 | 1 + 4 files changed, 3 insertions(+), 3 deletions(-) delete mode 100644 plugins/repository-hdfs/licenses/commons-lang3-3.15.0.jar.sha1 create mode 100644 plugins/repository-hdfs/licenses/commons-lang3-3.16.0.jar.sha1 diff --git a/CHANGELOG.md b/CHANGELOG.md index c60b8f492d02c..46402c0ab4ea4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -20,7 +20,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), ### Dependencies - Bump `netty` from 4.1.111.Final to 4.1.112.Final ([#15081](https://github.com/opensearch-project/OpenSearch/pull/15081)) -- Bump `org.apache.commons:commons-lang3` from 3.14.0 to 3.15.0 ([#14861](https://github.com/opensearch-project/OpenSearch/pull/14861)) +- Bump `org.apache.commons:commons-lang3` from 3.14.0 to 3.16.0 ([#14861](https://github.com/opensearch-project/OpenSearch/pull/14861), [#15205](https://github.com/opensearch-project/OpenSearch/pull/15205)) - OpenJDK Update (July 2024 Patch releases) ([#14998](https://github.com/opensearch-project/OpenSearch/pull/14998)) - Bump `com.microsoft.azure:msal4j` from 1.16.1 to 1.16.2 ([#14995](https://github.com/opensearch-project/OpenSearch/pull/14995)) - Bump `actions/github-script` from 6 to 7 ([#14997](https://github.com/opensearch-project/OpenSearch/pull/14997)) diff --git a/plugins/repository-hdfs/build.gradle b/plugins/repository-hdfs/build.gradle index f117bae658abe..e1e5f422f3e07 100644 --- a/plugins/repository-hdfs/build.gradle +++ b/plugins/repository-hdfs/build.gradle @@ -76,7 +76,7 @@ dependencies { api "org.apache.commons:commons-compress:${versions.commonscompress}" api 'org.apache.commons:commons-configuration2:2.11.0' api "commons-io:commons-io:${versions.commonsio}" - api 'org.apache.commons:commons-lang3:3.15.0' + api 'org.apache.commons:commons-lang3:3.16.0' implementation 'com.google.re2j:re2j:1.7' api 'javax.servlet:servlet-api:2.5' api "org.slf4j:slf4j-api:${versions.slf4j}" diff --git a/plugins/repository-hdfs/licenses/commons-lang3-3.15.0.jar.sha1 b/plugins/repository-hdfs/licenses/commons-lang3-3.15.0.jar.sha1 deleted file mode 100644 index 4b1179c935946..0000000000000 --- a/plugins/repository-hdfs/licenses/commons-lang3-3.15.0.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -21581109b4be710ea4b195d5760392ec284f9f11 \ No newline at end of file diff --git a/plugins/repository-hdfs/licenses/commons-lang3-3.16.0.jar.sha1 b/plugins/repository-hdfs/licenses/commons-lang3-3.16.0.jar.sha1 new file mode 100644 index 0000000000000..ef4f1c1fc2002 --- /dev/null +++ b/plugins/repository-hdfs/licenses/commons-lang3-3.16.0.jar.sha1 @@ -0,0 +1 @@ +3eb54effe40946dfb06dc5cd6c7ce4116cd51ea4 \ No newline at end of file From 23111954cdeb1f0430e2c87641dc8a24591b02d4 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Mon, 12 Aug 2024 12:15:49 -0400 Subject: [PATCH 21/40] Bump com.azure:azure-xml from 1.0.0 to 1.1.0 in /plugins/repository-azure (#15206) * Bump com.azure:azure-xml in /plugins/repository-azure Bumps [com.azure:azure-xml](https://github.com/Azure/azure-sdk-for-java) from 1.0.0 to 1.1.0. - [Release notes](https://github.com/Azure/azure-sdk-for-java/releases) - [Commits](https://github.com/Azure/azure-sdk-for-java/compare/v1.0.0...v1.1.0) --- updated-dependencies: - dependency-name: com.azure:azure-xml dependency-type: direct:production update-type: version-update:semver-minor ... Signed-off-by: dependabot[bot] * Updating SHAs Signed-off-by: dependabot[bot] * Update changelog Signed-off-by: dependabot[bot] --------- Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> Co-authored-by: dependabot[bot] --- CHANGELOG.md | 1 + plugins/repository-azure/build.gradle | 2 +- plugins/repository-azure/licenses/azure-xml-1.0.0.jar.sha1 | 1 - plugins/repository-azure/licenses/azure-xml-1.1.0.jar.sha1 | 1 + 4 files changed, 3 insertions(+), 2 deletions(-) delete mode 100644 plugins/repository-azure/licenses/azure-xml-1.0.0.jar.sha1 create mode 100644 plugins/repository-azure/licenses/azure-xml-1.1.0.jar.sha1 diff --git a/CHANGELOG.md b/CHANGELOG.md index 46402c0ab4ea4..8bb9799a36339 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -30,6 +30,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), - Bump `org.bouncycastle:bcpg-fips` from 1.0.7.1 to 2.0.8 and `org.bouncycastle:bc-fips` from 1.0.2.5 to 2.0.0 in /distribution/tools/plugin-cli ([#15103](https://github.com/opensearch-project/OpenSearch/pull/15103)) - Bump `com.azure:azure-core` from 1.49.1 to 1.51.0 ([#15111](https://github.com/opensearch-project/OpenSearch/pull/15111)) - Bump `org.xerial.snappy:snappy-java` from 1.1.10.5 to 1.1.10.6 ([#15207](https://github.com/opensearch-project/OpenSearch/pull/15207)) +- Bump `com.azure:azure-xml` from 1.0.0 to 1.1.0 ([#15206](https://github.com/opensearch-project/OpenSearch/pull/15206)) ### Changed - Add lower limit for primary and replica batch allocators timeout ([#14979](https://github.com/opensearch-project/OpenSearch/pull/14979)) diff --git a/plugins/repository-azure/build.gradle b/plugins/repository-azure/build.gradle index 80809e067f65a..6b63311cb3125 100644 --- a/plugins/repository-azure/build.gradle +++ b/plugins/repository-azure/build.gradle @@ -46,7 +46,7 @@ opensearchplugin { dependencies { api 'com.azure:azure-core:1.51.0' api 'com.azure:azure-json:1.1.0' - api 'com.azure:azure-xml:1.0.0' + api 'com.azure:azure-xml:1.1.0' api 'com.azure:azure-storage-common:12.25.1' api 'com.azure:azure-core-http-netty:1.15.1' api "io.netty:netty-codec-dns:${versions.netty}" diff --git a/plugins/repository-azure/licenses/azure-xml-1.0.0.jar.sha1 b/plugins/repository-azure/licenses/azure-xml-1.0.0.jar.sha1 deleted file mode 100644 index 798ec5d95c6ac..0000000000000 --- a/plugins/repository-azure/licenses/azure-xml-1.0.0.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -ba584703bd47e9e789343ee3332f0f5a64f7f187 \ No newline at end of file diff --git a/plugins/repository-azure/licenses/azure-xml-1.1.0.jar.sha1 b/plugins/repository-azure/licenses/azure-xml-1.1.0.jar.sha1 new file mode 100644 index 0000000000000..4f9cfcac02f6e --- /dev/null +++ b/plugins/repository-azure/licenses/azure-xml-1.1.0.jar.sha1 @@ -0,0 +1 @@ +8218a00c07f9f66d5dc7ae2ba613da6890867497 \ No newline at end of file From d442f7c1684af2c58e3de8c9b054c65072a03bea Mon Sep 17 00:00:00 2001 From: Ruirui Zhang Date: Mon, 12 Aug 2024 17:38:21 -0700 Subject: [PATCH 22/40] Add Create QueryGroup API Logic (#14680) * add logic for Create QueryGroup API Signed-off-by: Ruirui Zhang * remove wildcard imports Signed-off-by: Ruirui Zhang * change jvm to memeory Signed-off-by: Ruirui Zhang * modify querygroup Signed-off-by: Ruirui Zhang * fix javadoc and add more tests Signed-off-by: Ruirui Zhang * add more tests Signed-off-by: Ruirui Zhang * address comments Signed-off-by: Ruirui Zhang * fix the persist logic Signed-off-by: Ruirui Zhang * remove inflight checks as they are not necessary Signed-off-by: Ruirui Zhang * remove persistable interface Signed-off-by: Ruirui Zhang * modify QueryGroupServiceSettings Signed-off-by: Ruirui Zhang * add in an action package in the plugin Signed-off-by: Ruirui Zhang * modify based on commments Signed-off-by: Ruirui Zhang * address comments on QueryGroupPersistenceService Signed-off-by: Ruirui Zhang * address comments on persistence service Signed-off-by: Ruirui Zhang * address comments Signed-off-by: Ruirui Zhang * fix unit test Signed-off-by: Ruirui Zhang * address comments Signed-off-by: Ruirui Zhang * add IT Signed-off-by: Ruirui Zhang * add coverage Signed-off-by: Ruirui Zhang --- CHANGELOG.md | 1 + plugins/workload-management/build.gradle | 21 ++ .../plugin/wlm/WorkloadManagementPlugin.java | 64 +++++ .../wlm/action/CreateQueryGroupAction.java | 36 +++ .../wlm/action/CreateQueryGroupRequest.java | 82 ++++++ .../wlm/action/CreateQueryGroupResponse.java | 74 ++++++ .../TransportCreateQueryGroupAction.java | 57 ++++ .../plugin/wlm/action/package-info.java | 12 + .../opensearch/plugin/wlm/package-info.java | 12 + .../wlm/rest/RestCreateQueryGroupAction.java | 72 +++++ .../plugin/wlm/rest/package-info.java | 12 + .../service/QueryGroupPersistenceService.java | 201 ++++++++++++++ .../plugin/wlm/service/package-info.java | 12 + .../plugin/wlm/QueryGroupTestUtils.java | 141 ++++++++++ .../action/CreateQueryGroupRequestTests.java | 40 +++ .../action/CreateQueryGroupResponseTests.java | 66 +++++ .../QueryGroupPersistenceServiceTests.java | 247 ++++++++++++++++++ ...rkloadManagementClientYamlTestSuiteIT.java | 52 ++++ .../api/create_query_group_context.json | 18 ++ .../test/wlm/10_create_query_group.yml | 90 +++++++ .../opensearch/cluster/metadata/Metadata.java | 6 + .../cluster/metadata/QueryGroup.java | 157 +++++------ .../metadata/QueryGroupMetadataTests.java | 2 +- .../cluster/metadata/QueryGroupTests.java | 33 ++- 24 files changed, 1427 insertions(+), 81 deletions(-) create mode 100644 plugins/workload-management/build.gradle create mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java create mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupAction.java create mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequest.java create mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupResponse.java create mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportCreateQueryGroupAction.java create mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/package-info.java create mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/package-info.java create mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestCreateQueryGroupAction.java create mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/package-info.java create mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceService.java create mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/package-info.java create mode 100644 plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/QueryGroupTestUtils.java create mode 100644 plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequestTests.java create mode 100644 plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateQueryGroupResponseTests.java create mode 100644 plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceServiceTests.java create mode 100644 plugins/workload-management/src/yamlRestTest/java/org/opensearch/plugin/wlm/WorkloadManagementClientYamlTestSuiteIT.java create mode 100644 plugins/workload-management/src/yamlRestTest/resources/rest-api-spec/api/create_query_group_context.json create mode 100644 plugins/workload-management/src/yamlRestTest/resources/rest-api-spec/test/wlm/10_create_query_group.yml diff --git a/CHANGELOG.md b/CHANGELOG.md index 8bb9799a36339..34cd4c2097e48 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), - Add setting to ignore throttling nodes for allocation of unassigned primaries in remote restore ([#14991](https://github.com/opensearch-project/OpenSearch/pull/14991)) - [Streaming Indexing] Enhance RestClient with a new streaming API support ([#14437](https://github.com/opensearch-project/OpenSearch/pull/14437)) - Add basic aggregation support for derived fields ([#14618](https://github.com/opensearch-project/OpenSearch/pull/14618)) +- [Workload Management] Add Create QueryGroup API Logic ([#14680](https://github.com/opensearch-project/OpenSearch/pull/14680))- [Workload Management] Add Create QueryGroup API Logic ([#14680](https://github.com/opensearch-project/OpenSearch/pull/14680)) - Add ThreadContextPermission for markAsSystemContext and allow core to perform the method ([#15016](https://github.com/opensearch-project/OpenSearch/pull/15016)) - Add ThreadContextPermission for stashAndMergeHeaders and stashWithOrigin ([#15039](https://github.com/opensearch-project/OpenSearch/pull/15039)) - [Concurrent Segment Search] Support composite aggregations with scripting ([#15072](https://github.com/opensearch-project/OpenSearch/pull/15072)) diff --git a/plugins/workload-management/build.gradle b/plugins/workload-management/build.gradle new file mode 100644 index 0000000000000..cb14d22ef149f --- /dev/null +++ b/plugins/workload-management/build.gradle @@ -0,0 +1,21 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + * + * Modifications Copyright OpenSearch Contributors. See + * GitHub history for details. + */ + +apply plugin: 'opensearch.yaml-rest-test' +apply plugin: 'opensearch.internal-cluster-test' + +opensearchplugin { + description 'OpenSearch Workload Management Plugin.' + classname 'org.opensearch.plugin.wlm.WorkloadManagementPlugin' +} + +dependencies { +} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java new file mode 100644 index 0000000000000..80807f0d5bc37 --- /dev/null +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java @@ -0,0 +1,64 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.plugin.wlm; + +import org.opensearch.action.ActionRequest; +import org.opensearch.cluster.metadata.IndexNameExpressionResolver; +import org.opensearch.cluster.node.DiscoveryNodes; +import org.opensearch.common.settings.ClusterSettings; +import org.opensearch.common.settings.IndexScopedSettings; +import org.opensearch.common.settings.Setting; +import org.opensearch.common.settings.Settings; +import org.opensearch.common.settings.SettingsFilter; +import org.opensearch.core.action.ActionResponse; +import org.opensearch.plugin.wlm.action.CreateQueryGroupAction; +import org.opensearch.plugin.wlm.action.TransportCreateQueryGroupAction; +import org.opensearch.plugin.wlm.rest.RestCreateQueryGroupAction; +import org.opensearch.plugin.wlm.service.QueryGroupPersistenceService; +import org.opensearch.plugins.ActionPlugin; +import org.opensearch.plugins.Plugin; +import org.opensearch.rest.RestController; +import org.opensearch.rest.RestHandler; + +import java.util.List; +import java.util.function.Supplier; + +/** + * Plugin class for WorkloadManagement + */ +public class WorkloadManagementPlugin extends Plugin implements ActionPlugin { + + /** + * Default constructor + */ + public WorkloadManagementPlugin() {} + + @Override + public List> getActions() { + return List.of(new ActionPlugin.ActionHandler<>(CreateQueryGroupAction.INSTANCE, TransportCreateQueryGroupAction.class)); + } + + @Override + public List getRestHandlers( + Settings settings, + RestController restController, + ClusterSettings clusterSettings, + IndexScopedSettings indexScopedSettings, + SettingsFilter settingsFilter, + IndexNameExpressionResolver indexNameExpressionResolver, + Supplier nodesInCluster + ) { + return List.of(new RestCreateQueryGroupAction()); + } + + @Override + public List> getSettings() { + return List.of(QueryGroupPersistenceService.MAX_QUERY_GROUP_COUNT); + } +} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupAction.java new file mode 100644 index 0000000000000..14cb8cfcd125a --- /dev/null +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupAction.java @@ -0,0 +1,36 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.plugin.wlm.action; + +import org.opensearch.action.ActionType; + +/** + * Transport action to create QueryGroup + * + * @opensearch.experimental + */ +public class CreateQueryGroupAction extends ActionType { + + /** + * An instance of CreateQueryGroupAction + */ + public static final CreateQueryGroupAction INSTANCE = new CreateQueryGroupAction(); + + /** + * Name for CreateQueryGroupAction + */ + public static final String NAME = "cluster:admin/opensearch/wlm/query_group/_create"; + + /** + * Default constructor + */ + private CreateQueryGroupAction() { + super(NAME, CreateQueryGroupResponse::new); + } +} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequest.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequest.java new file mode 100644 index 0000000000000..ff6422be36885 --- /dev/null +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequest.java @@ -0,0 +1,82 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.plugin.wlm.action; + +import org.opensearch.action.ActionRequest; +import org.opensearch.action.ActionRequestValidationException; +import org.opensearch.cluster.metadata.QueryGroup; +import org.opensearch.common.UUIDs; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.core.common.io.stream.StreamOutput; +import org.opensearch.core.xcontent.XContentParser; +import org.joda.time.Instant; + +import java.io.IOException; + +/** + * A request for create QueryGroup + * User input schema: + * { + * "name": "analytics", + * "resiliency_mode": "enforced", + * "resource_limits": { + * "cpu" : 0.4, + * "memory" : 0.2 + * } + * } + * + * @opensearch.experimental + */ +public class CreateQueryGroupRequest extends ActionRequest { + private final QueryGroup queryGroup; + + /** + * Constructor for CreateQueryGroupRequest + * @param queryGroup - A {@link QueryGroup} object + */ + public CreateQueryGroupRequest(QueryGroup queryGroup) { + this.queryGroup = queryGroup; + } + + /** + * Constructor for CreateQueryGroupRequest + * @param in - A {@link StreamInput} object + */ + public CreateQueryGroupRequest(StreamInput in) throws IOException { + super(in); + queryGroup = new QueryGroup(in); + } + + /** + * Generate a CreateQueryGroupRequest from XContent + * @param parser - A {@link XContentParser} object + */ + public static CreateQueryGroupRequest fromXContent(XContentParser parser) throws IOException { + QueryGroup.Builder builder = QueryGroup.Builder.fromXContent(parser); + return new CreateQueryGroupRequest(builder._id(UUIDs.randomBase64UUID()).updatedAt(Instant.now().getMillis()).build()); + } + + @Override + public ActionRequestValidationException validate() { + return null; + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + super.writeTo(out); + queryGroup.writeTo(out); + } + + /** + * QueryGroup getter + */ + public QueryGroup getQueryGroup() { + return queryGroup; + } +} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupResponse.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupResponse.java new file mode 100644 index 0000000000000..9a2a8178c0a29 --- /dev/null +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupResponse.java @@ -0,0 +1,74 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.plugin.wlm.action; + +import org.opensearch.cluster.metadata.QueryGroup; +import org.opensearch.core.action.ActionResponse; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.core.common.io.stream.StreamOutput; +import org.opensearch.core.rest.RestStatus; +import org.opensearch.core.xcontent.ToXContent; +import org.opensearch.core.xcontent.ToXContentObject; +import org.opensearch.core.xcontent.XContentBuilder; + +import java.io.IOException; + +/** + * Response for the create API for QueryGroup + * + * @opensearch.experimental + */ +public class CreateQueryGroupResponse extends ActionResponse implements ToXContent, ToXContentObject { + private final QueryGroup queryGroup; + private final RestStatus restStatus; + + /** + * Constructor for CreateQueryGroupResponse + * @param queryGroup - The QueryGroup to be included in the response + * @param restStatus - The restStatus for the response + */ + public CreateQueryGroupResponse(final QueryGroup queryGroup, RestStatus restStatus) { + this.queryGroup = queryGroup; + this.restStatus = restStatus; + } + + /** + * Constructor for CreateQueryGroupResponse + * @param in - A {@link StreamInput} object + */ + public CreateQueryGroupResponse(StreamInput in) throws IOException { + queryGroup = new QueryGroup(in); + restStatus = RestStatus.readFrom(in); + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + queryGroup.writeTo(out); + RestStatus.writeTo(out, restStatus); + } + + @Override + public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { + return queryGroup.toXContent(builder, params); + } + + /** + * queryGroup getter + */ + public QueryGroup getQueryGroup() { + return queryGroup; + } + + /** + * restStatus getter + */ + public RestStatus getRestStatus() { + return restStatus; + } +} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportCreateQueryGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportCreateQueryGroupAction.java new file mode 100644 index 0000000000000..01aa8cfb5e610 --- /dev/null +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportCreateQueryGroupAction.java @@ -0,0 +1,57 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.plugin.wlm.action; + +import org.opensearch.action.support.ActionFilters; +import org.opensearch.action.support.HandledTransportAction; +import org.opensearch.common.inject.Inject; +import org.opensearch.core.action.ActionListener; +import org.opensearch.plugin.wlm.service.QueryGroupPersistenceService; +import org.opensearch.tasks.Task; +import org.opensearch.threadpool.ThreadPool; +import org.opensearch.transport.TransportService; + +/** + * Transport action to create QueryGroup + * + * @opensearch.experimental + */ +public class TransportCreateQueryGroupAction extends HandledTransportAction { + + private final ThreadPool threadPool; + private final QueryGroupPersistenceService queryGroupPersistenceService; + + /** + * Constructor for TransportCreateQueryGroupAction + * + * @param actionName - action name + * @param transportService - a {@link TransportService} object + * @param actionFilters - a {@link ActionFilters} object + * @param threadPool - a {@link ThreadPool} object + * @param queryGroupPersistenceService - a {@link QueryGroupPersistenceService} object + */ + @Inject + public TransportCreateQueryGroupAction( + String actionName, + TransportService transportService, + ActionFilters actionFilters, + ThreadPool threadPool, + QueryGroupPersistenceService queryGroupPersistenceService + ) { + super(CreateQueryGroupAction.NAME, transportService, actionFilters, CreateQueryGroupRequest::new); + this.threadPool = threadPool; + this.queryGroupPersistenceService = queryGroupPersistenceService; + } + + @Override + protected void doExecute(Task task, CreateQueryGroupRequest request, ActionListener listener) { + threadPool.executor(ThreadPool.Names.SAME) + .execute(() -> queryGroupPersistenceService.persistInClusterStateMetadata(request.getQueryGroup(), listener)); + } +} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/package-info.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/package-info.java new file mode 100644 index 0000000000000..9921500df8a81 --- /dev/null +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/package-info.java @@ -0,0 +1,12 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +/** + * Package for the action classes of WorkloadManagementPlugin + */ +package org.opensearch.plugin.wlm.action; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/package-info.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/package-info.java new file mode 100644 index 0000000000000..84c99967b226b --- /dev/null +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/package-info.java @@ -0,0 +1,12 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +/** + * Base package for WorkloadManagementPlugin + */ +package org.opensearch.plugin.wlm; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestCreateQueryGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestCreateQueryGroupAction.java new file mode 100644 index 0000000000000..b0e0af4f9d17f --- /dev/null +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestCreateQueryGroupAction.java @@ -0,0 +1,72 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.plugin.wlm.rest; + +import org.opensearch.client.node.NodeClient; +import org.opensearch.core.rest.RestStatus; +import org.opensearch.core.xcontent.ToXContent; +import org.opensearch.core.xcontent.XContentParser; +import org.opensearch.plugin.wlm.action.CreateQueryGroupAction; +import org.opensearch.plugin.wlm.action.CreateQueryGroupRequest; +import org.opensearch.plugin.wlm.action.CreateQueryGroupResponse; +import org.opensearch.rest.BaseRestHandler; +import org.opensearch.rest.BytesRestResponse; +import org.opensearch.rest.RestChannel; +import org.opensearch.rest.RestRequest; +import org.opensearch.rest.RestResponse; +import org.opensearch.rest.action.RestResponseListener; + +import java.io.IOException; +import java.util.List; + +import static org.opensearch.rest.RestRequest.Method.POST; +import static org.opensearch.rest.RestRequest.Method.PUT; + +/** + * Rest action to create a QueryGroup + * + * @opensearch.experimental + */ +public class RestCreateQueryGroupAction extends BaseRestHandler { + + /** + * Constructor for RestCreateQueryGroupAction + */ + public RestCreateQueryGroupAction() {} + + @Override + public String getName() { + return "create_query_group"; + } + + /** + * The list of {@link Route}s that this RestHandler is responsible for handling. + */ + @Override + public List routes() { + return List.of(new Route(POST, "_wlm/query_group/"), new Route(PUT, "_wlm/query_group/")); + } + + @Override + protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException { + try (XContentParser parser = request.contentParser()) { + CreateQueryGroupRequest createQueryGroupRequest = CreateQueryGroupRequest.fromXContent(parser); + return channel -> client.execute(CreateQueryGroupAction.INSTANCE, createQueryGroupRequest, createQueryGroupResponse(channel)); + } + } + + private RestResponseListener createQueryGroupResponse(final RestChannel channel) { + return new RestResponseListener<>(channel) { + @Override + public RestResponse buildResponse(final CreateQueryGroupResponse response) throws Exception { + return new BytesRestResponse(RestStatus.OK, response.toXContent(channel.newBuilder(), ToXContent.EMPTY_PARAMS)); + } + }; + } +} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/package-info.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/package-info.java new file mode 100644 index 0000000000000..7d7cb9028fdb8 --- /dev/null +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/package-info.java @@ -0,0 +1,12 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +/** + * Package for the rest classes of WorkloadManagementPlugin + */ +package org.opensearch.plugin.wlm.rest; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceService.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceService.java new file mode 100644 index 0000000000000..b2164df561bf9 --- /dev/null +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceService.java @@ -0,0 +1,201 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.plugin.wlm.service; + +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.opensearch.cluster.ClusterState; +import org.opensearch.cluster.ClusterStateUpdateTask; +import org.opensearch.cluster.metadata.Metadata; +import org.opensearch.cluster.metadata.QueryGroup; +import org.opensearch.cluster.service.ClusterManagerTaskThrottler.ThrottlingKey; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.Priority; +import org.opensearch.common.inject.Inject; +import org.opensearch.common.settings.ClusterSettings; +import org.opensearch.common.settings.Setting; +import org.opensearch.common.settings.Settings; +import org.opensearch.core.action.ActionListener; +import org.opensearch.core.rest.RestStatus; +import org.opensearch.plugin.wlm.action.CreateQueryGroupResponse; +import org.opensearch.search.ResourceType; + +import java.util.EnumMap; +import java.util.Map; +import java.util.Optional; + +/** + * This class defines the functions for QueryGroup persistence + */ +public class QueryGroupPersistenceService { + static final String SOURCE = "query-group-persistence-service"; + private static final String CREATE_QUERY_GROUP_THROTTLING_KEY = "create-query-group"; + private static final Logger logger = LogManager.getLogger(QueryGroupPersistenceService.class); + /** + * max QueryGroup count setting name + */ + public static final String QUERY_GROUP_COUNT_SETTING_NAME = "node.query_group.max_count"; + /** + * default max queryGroup count on any node at any given point in time + */ + private static final int DEFAULT_MAX_QUERY_GROUP_COUNT_VALUE = 100; + /** + * min queryGroup count on any node at any given point in time + */ + private static final int MIN_QUERY_GROUP_COUNT_VALUE = 1; + /** + * max QueryGroup count setting + */ + public static final Setting MAX_QUERY_GROUP_COUNT = Setting.intSetting( + QUERY_GROUP_COUNT_SETTING_NAME, + DEFAULT_MAX_QUERY_GROUP_COUNT_VALUE, + 0, + QueryGroupPersistenceService::validateMaxQueryGroupCount, + Setting.Property.Dynamic, + Setting.Property.NodeScope + ); + private final ClusterService clusterService; + private volatile int maxQueryGroupCount; + final ThrottlingKey createQueryGroupThrottlingKey; + + /** + * Constructor for QueryGroupPersistenceService + * + * @param clusterService {@link ClusterService} - The cluster service to be used by QueryGroupPersistenceService + * @param settings {@link Settings} - The settings to be used by QueryGroupPersistenceService + * @param clusterSettings {@link ClusterSettings} - The cluster settings to be used by QueryGroupPersistenceService + */ + @Inject + public QueryGroupPersistenceService( + final ClusterService clusterService, + final Settings settings, + final ClusterSettings clusterSettings + ) { + this.clusterService = clusterService; + this.createQueryGroupThrottlingKey = clusterService.registerClusterManagerTask(CREATE_QUERY_GROUP_THROTTLING_KEY, true); + setMaxQueryGroupCount(MAX_QUERY_GROUP_COUNT.get(settings)); + clusterSettings.addSettingsUpdateConsumer(MAX_QUERY_GROUP_COUNT, this::setMaxQueryGroupCount); + } + + /** + * Set maxQueryGroupCount to be newMaxQueryGroupCount + * @param newMaxQueryGroupCount - the max number of QueryGroup allowed + */ + public void setMaxQueryGroupCount(int newMaxQueryGroupCount) { + validateMaxQueryGroupCount(newMaxQueryGroupCount); + this.maxQueryGroupCount = newMaxQueryGroupCount; + } + + /** + * Validator for maxQueryGroupCount + * @param maxQueryGroupCount - the maxQueryGroupCount number to be verified + */ + private static void validateMaxQueryGroupCount(int maxQueryGroupCount) { + if (maxQueryGroupCount > DEFAULT_MAX_QUERY_GROUP_COUNT_VALUE || maxQueryGroupCount < MIN_QUERY_GROUP_COUNT_VALUE) { + throw new IllegalArgumentException(QUERY_GROUP_COUNT_SETTING_NAME + " should be in range [1-100]."); + } + } + + /** + * Update cluster state to include the new QueryGroup + * @param queryGroup {@link QueryGroup} - the QueryGroup we're currently creating + * @param listener - ActionListener for CreateQueryGroupResponse + */ + public void persistInClusterStateMetadata(QueryGroup queryGroup, ActionListener listener) { + clusterService.submitStateUpdateTask(SOURCE, new ClusterStateUpdateTask(Priority.NORMAL) { + @Override + public ClusterState execute(ClusterState currentState) throws Exception { + return saveQueryGroupInClusterState(queryGroup, currentState); + } + + @Override + public ThrottlingKey getClusterManagerThrottlingKey() { + return createQueryGroupThrottlingKey; + } + + @Override + public void onFailure(String source, Exception e) { + logger.warn("failed to save QueryGroup object due to error: {}, for source: {}.", e.getMessage(), source); + listener.onFailure(e); + } + + @Override + public void clusterStateProcessed(String source, ClusterState oldState, ClusterState newState) { + CreateQueryGroupResponse response = new CreateQueryGroupResponse(queryGroup, RestStatus.OK); + listener.onResponse(response); + } + }); + } + + /** + * This method will be executed before we submit the new cluster state + * @param queryGroup - the QueryGroup we're currently creating + * @param currentClusterState - the cluster state before the update + */ + ClusterState saveQueryGroupInClusterState(final QueryGroup queryGroup, final ClusterState currentClusterState) { + final Map existingQueryGroups = currentClusterState.metadata().queryGroups(); + String groupName = queryGroup.getName(); + + // check if maxQueryGroupCount will breach + if (existingQueryGroups.size() == maxQueryGroupCount) { + logger.warn("{} value exceeded its assigned limit of {}.", QUERY_GROUP_COUNT_SETTING_NAME, maxQueryGroupCount); + throw new IllegalStateException("Can't create more than " + maxQueryGroupCount + " QueryGroups in the system."); + } + + // check for duplicate name + Optional findExistingGroup = existingQueryGroups.values() + .stream() + .filter(group -> group.getName().equals(groupName)) + .findFirst(); + if (findExistingGroup.isPresent()) { + logger.warn("QueryGroup with name {} already exists. Not creating a new one.", groupName); + throw new IllegalArgumentException("QueryGroup with name " + groupName + " already exists. Not creating a new one."); + } + + // check if there's any resource allocation that exceed limit of 1.0 + Map totalUsageMap = calculateTotalUsage(existingQueryGroups, queryGroup); + for (ResourceType resourceType : queryGroup.getResourceLimits().keySet()) { + if (totalUsageMap.get(resourceType) > 1) { + logger.warn("Total resource allocation for {} will go above the max limit of 1.0.", resourceType.getName()); + throw new IllegalArgumentException( + "Total resource allocation for " + resourceType.getName() + " will go above the max limit of 1.0." + ); + } + } + + return ClusterState.builder(currentClusterState) + .metadata(Metadata.builder(currentClusterState.metadata()).put(queryGroup).build()) + .build(); + } + + /** + * This method calculates the existing total usage of the all the resource limits + * @param existingQueryGroups - existing QueryGroups in the system + * @param queryGroup - the QueryGroup we're creating or updating + */ + private Map calculateTotalUsage(Map existingQueryGroups, QueryGroup queryGroup) { + final Map map = new EnumMap<>(ResourceType.class); + map.putAll(queryGroup.getResourceLimits()); + for (QueryGroup currGroup : existingQueryGroups.values()) { + if (!currGroup.getName().equals(queryGroup.getName())) { + for (ResourceType resourceType : queryGroup.getResourceLimits().keySet()) { + map.compute(resourceType, (k, v) -> v + currGroup.getResourceLimits().get(resourceType)); + } + } + } + return map; + } + + /** + * maxQueryGroupCount getter + */ + public int getMaxQueryGroupCount() { + return maxQueryGroupCount; + } +} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/package-info.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/package-info.java new file mode 100644 index 0000000000000..5848e9c936623 --- /dev/null +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/package-info.java @@ -0,0 +1,12 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +/** + * Package for the service classes of WorkloadManagementPlugin + */ +package org.opensearch.plugin.wlm.service; diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/QueryGroupTestUtils.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/QueryGroupTestUtils.java new file mode 100644 index 0000000000000..fc324853d9b34 --- /dev/null +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/QueryGroupTestUtils.java @@ -0,0 +1,141 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.plugin.wlm; + +import org.opensearch.cluster.ClusterName; +import org.opensearch.cluster.ClusterState; +import org.opensearch.cluster.metadata.Metadata; +import org.opensearch.cluster.metadata.QueryGroup; +import org.opensearch.cluster.service.ClusterApplierService; +import org.opensearch.cluster.service.ClusterManagerService; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.collect.Tuple; +import org.opensearch.common.settings.ClusterSettings; +import org.opensearch.common.settings.Setting; +import org.opensearch.common.settings.Settings; +import org.opensearch.plugin.wlm.service.QueryGroupPersistenceService; +import org.opensearch.threadpool.ThreadPool; + +import java.util.ArrayList; +import java.util.Comparator; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; + +import static org.opensearch.cluster.metadata.QueryGroup.builder; +import static org.opensearch.search.ResourceType.fromName; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; +import static org.mockito.Mockito.mock; + +public class QueryGroupTestUtils { + public static final String NAME_ONE = "query_group_one"; + public static final String NAME_TWO = "query_group_two"; + public static final String _ID_ONE = "AgfUO5Ja9yfsYlONlYi3TQ=="; + public static final String _ID_TWO = "G5iIqHy4g7eK1qIAAAAIH53=1"; + public static final String NAME_NONE_EXISTED = "query_group_none_existed"; + public static final String MEMORY_STRING = "memory"; + public static final String MONITOR_STRING = "monitor"; + public static final long TIMESTAMP_ONE = 4513232413L; + public static final long TIMESTAMP_TWO = 4513232415L; + public static final QueryGroup queryGroupOne = builder().name(NAME_ONE) + ._id(_ID_ONE) + .mode(MONITOR_STRING) + .resourceLimits(Map.of(fromName(MEMORY_STRING), 0.3)) + .updatedAt(TIMESTAMP_ONE) + .build(); + + public static final QueryGroup queryGroupTwo = builder().name(NAME_TWO) + ._id(_ID_TWO) + .mode(MONITOR_STRING) + .resourceLimits(Map.of(fromName(MEMORY_STRING), 0.6)) + .updatedAt(TIMESTAMP_TWO) + .build(); + + public static List queryGroupList() { + List list = new ArrayList<>(); + list.add(queryGroupOne); + list.add(queryGroupTwo); + return list; + } + + public static ClusterState clusterState() { + final Metadata metadata = Metadata.builder().queryGroups(Map.of(_ID_ONE, queryGroupOne, _ID_TWO, queryGroupTwo)).build(); + return ClusterState.builder(new ClusterName("_name")).metadata(metadata).build(); + } + + public static Set> clusterSettingsSet() { + Set> set = new HashSet<>(ClusterSettings.BUILT_IN_CLUSTER_SETTINGS); + set.add(QueryGroupPersistenceService.MAX_QUERY_GROUP_COUNT); + assertFalse(ClusterSettings.BUILT_IN_CLUSTER_SETTINGS.contains(QueryGroupPersistenceService.MAX_QUERY_GROUP_COUNT)); + return set; + } + + public static Settings settings() { + return Settings.builder().build(); + } + + public static ClusterSettings clusterSettings() { + return new ClusterSettings(settings(), clusterSettingsSet()); + } + + public static QueryGroupPersistenceService queryGroupPersistenceService() { + ClusterApplierService clusterApplierService = new ClusterApplierService( + "name", + settings(), + clusterSettings(), + mock(ThreadPool.class) + ); + clusterApplierService.setInitialState(clusterState()); + ClusterService clusterService = new ClusterService( + settings(), + clusterSettings(), + mock(ClusterManagerService.class), + clusterApplierService + ); + return new QueryGroupPersistenceService(clusterService, settings(), clusterSettings()); + } + + public static Tuple preparePersistenceServiceSetup(Map queryGroups) { + Metadata metadata = Metadata.builder().queryGroups(queryGroups).build(); + Settings settings = Settings.builder().build(); + ClusterState clusterState = ClusterState.builder(new ClusterName("_name")).metadata(metadata).build(); + ClusterSettings clusterSettings = new ClusterSettings(settings, clusterSettingsSet()); + ClusterApplierService clusterApplierService = new ClusterApplierService( + "name", + settings(), + clusterSettings(), + mock(ThreadPool.class) + ); + clusterApplierService.setInitialState(clusterState); + ClusterService clusterService = new ClusterService( + settings(), + clusterSettings(), + mock(ClusterManagerService.class), + clusterApplierService + ); + QueryGroupPersistenceService queryGroupPersistenceService = new QueryGroupPersistenceService( + clusterService, + settings, + clusterSettings + ); + return new Tuple(queryGroupPersistenceService, clusterState); + } + + public static void assertEqualQueryGroups(List listOne, List listTwo) { + assertEquals(listOne.size(), listTwo.size()); + listOne.sort(Comparator.comparing(QueryGroup::getName)); + listTwo.sort(Comparator.comparing(QueryGroup::getName)); + for (int i = 0; i < listOne.size(); i++) { + assertTrue(listOne.get(i).equals(listTwo.get(i))); + } + } +} diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequestTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequestTests.java new file mode 100644 index 0000000000000..b0fa96a46df80 --- /dev/null +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequestTests.java @@ -0,0 +1,40 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.plugin.wlm.action; + +import org.opensearch.cluster.metadata.QueryGroup; +import org.opensearch.common.io.stream.BytesStreamOutput; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.test.OpenSearchTestCase; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; + +import static org.opensearch.plugin.wlm.QueryGroupTestUtils.assertEqualQueryGroups; +import static org.opensearch.plugin.wlm.QueryGroupTestUtils.queryGroupOne; + +public class CreateQueryGroupRequestTests extends OpenSearchTestCase { + + /** + * Test case to verify the serialization and deserialization of CreateQueryGroupRequest. + */ + public void testSerialization() throws IOException { + CreateQueryGroupRequest request = new CreateQueryGroupRequest(queryGroupOne); + BytesStreamOutput out = new BytesStreamOutput(); + request.writeTo(out); + StreamInput streamInput = out.bytes().streamInput(); + CreateQueryGroupRequest otherRequest = new CreateQueryGroupRequest(streamInput); + List list1 = new ArrayList<>(); + List list2 = new ArrayList<>(); + list1.add(queryGroupOne); + list2.add(otherRequest.getQueryGroup()); + assertEqualQueryGroups(list1, list2); + } +} diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateQueryGroupResponseTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateQueryGroupResponseTests.java new file mode 100644 index 0000000000000..038f015713c5b --- /dev/null +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateQueryGroupResponseTests.java @@ -0,0 +1,66 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.plugin.wlm.action; + +import org.opensearch.cluster.metadata.QueryGroup; +import org.opensearch.common.io.stream.BytesStreamOutput; +import org.opensearch.common.xcontent.json.JsonXContent; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.core.rest.RestStatus; +import org.opensearch.core.xcontent.ToXContent; +import org.opensearch.core.xcontent.XContentBuilder; +import org.opensearch.plugin.wlm.QueryGroupTestUtils; +import org.opensearch.test.OpenSearchTestCase; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; + +import static org.mockito.Mockito.mock; + +public class CreateQueryGroupResponseTests extends OpenSearchTestCase { + + /** + * Test case to verify the serialization and deserialization of CreateQueryGroupResponse. + */ + public void testSerialization() throws IOException { + CreateQueryGroupResponse response = new CreateQueryGroupResponse(QueryGroupTestUtils.queryGroupOne, RestStatus.OK); + BytesStreamOutput out = new BytesStreamOutput(); + response.writeTo(out); + StreamInput streamInput = out.bytes().streamInput(); + CreateQueryGroupResponse otherResponse = new CreateQueryGroupResponse(streamInput); + assertEquals(response.getRestStatus(), otherResponse.getRestStatus()); + QueryGroup responseGroup = response.getQueryGroup(); + QueryGroup otherResponseGroup = otherResponse.getQueryGroup(); + List listOne = new ArrayList<>(); + List listTwo = new ArrayList<>(); + listOne.add(responseGroup); + listTwo.add(otherResponseGroup); + QueryGroupTestUtils.assertEqualQueryGroups(listOne, listTwo); + } + + /** + * Test case to verify the toXContent method of CreateQueryGroupResponse. + */ + public void testToXContentCreateQueryGroup() throws IOException { + XContentBuilder builder = JsonXContent.contentBuilder().prettyPrint(); + CreateQueryGroupResponse response = new CreateQueryGroupResponse(QueryGroupTestUtils.queryGroupOne, RestStatus.OK); + String actual = response.toXContent(builder, mock(ToXContent.Params.class)).toString(); + String expected = "{\n" + + " \"_id\" : \"AgfUO5Ja9yfsYlONlYi3TQ==\",\n" + + " \"name\" : \"query_group_one\",\n" + + " \"resiliency_mode\" : \"monitor\",\n" + + " \"updated_at\" : 4513232413,\n" + + " \"resource_limits\" : {\n" + + " \"memory\" : 0.3\n" + + " }\n" + + "}"; + assertEquals(expected, actual); + } +} diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceServiceTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceServiceTests.java new file mode 100644 index 0000000000000..533c98b44685d --- /dev/null +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceServiceTests.java @@ -0,0 +1,247 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.plugin.wlm.service; + +import org.opensearch.cluster.ClusterName; +import org.opensearch.cluster.ClusterState; +import org.opensearch.cluster.ClusterStateUpdateTask; +import org.opensearch.cluster.metadata.Metadata; +import org.opensearch.cluster.metadata.QueryGroup; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.collect.Tuple; +import org.opensearch.common.settings.ClusterSettings; +import org.opensearch.common.settings.Settings; +import org.opensearch.core.action.ActionListener; +import org.opensearch.plugin.wlm.QueryGroupTestUtils; +import org.opensearch.plugin.wlm.action.CreateQueryGroupResponse; +import org.opensearch.search.ResourceType; +import org.opensearch.test.OpenSearchTestCase; +import org.opensearch.threadpool.ThreadPool; + +import java.util.ArrayList; +import java.util.Collection; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +import org.mockito.ArgumentCaptor; + +import static org.opensearch.cluster.metadata.QueryGroup.builder; +import static org.opensearch.plugin.wlm.QueryGroupTestUtils.MEMORY_STRING; +import static org.opensearch.plugin.wlm.QueryGroupTestUtils.MONITOR_STRING; +import static org.opensearch.plugin.wlm.QueryGroupTestUtils.NAME_NONE_EXISTED; +import static org.opensearch.plugin.wlm.QueryGroupTestUtils.NAME_ONE; +import static org.opensearch.plugin.wlm.QueryGroupTestUtils._ID_ONE; +import static org.opensearch.plugin.wlm.QueryGroupTestUtils._ID_TWO; +import static org.opensearch.plugin.wlm.QueryGroupTestUtils.assertEqualQueryGroups; +import static org.opensearch.plugin.wlm.QueryGroupTestUtils.clusterSettings; +import static org.opensearch.plugin.wlm.QueryGroupTestUtils.clusterSettingsSet; +import static org.opensearch.plugin.wlm.QueryGroupTestUtils.preparePersistenceServiceSetup; +import static org.opensearch.plugin.wlm.QueryGroupTestUtils.queryGroupList; +import static org.opensearch.plugin.wlm.QueryGroupTestUtils.queryGroupOne; +import static org.opensearch.plugin.wlm.QueryGroupTestUtils.queryGroupTwo; +import static org.opensearch.plugin.wlm.service.QueryGroupPersistenceService.QUERY_GROUP_COUNT_SETTING_NAME; +import static org.opensearch.plugin.wlm.service.QueryGroupPersistenceService.SOURCE; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.anyString; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; + +public class QueryGroupPersistenceServiceTests extends OpenSearchTestCase { + + /** + * Test case to validate the creation logic of a QueryGroup + */ + public void testCreateQueryGroup() { + Tuple setup = preparePersistenceServiceSetup(new HashMap<>()); + QueryGroupPersistenceService queryGroupPersistenceService1 = setup.v1(); + ClusterState clusterState = setup.v2(); + ClusterState newClusterState = queryGroupPersistenceService1.saveQueryGroupInClusterState(queryGroupOne, clusterState); + Map updatedGroupsMap = newClusterState.getMetadata().queryGroups(); + assertEquals(1, updatedGroupsMap.size()); + assertTrue(updatedGroupsMap.containsKey(_ID_ONE)); + List listOne = new ArrayList<>(); + List listTwo = new ArrayList<>(); + listOne.add(queryGroupOne); + listTwo.add(updatedGroupsMap.get(_ID_ONE)); + assertEqualQueryGroups(listOne, listTwo); + } + + /** + * Test case to validate the logic for adding a new QueryGroup to a cluster state that already contains + * an existing QueryGroup + */ + public void testCreateAnotherQueryGroup() { + Tuple setup = preparePersistenceServiceSetup(Map.of(_ID_ONE, queryGroupOne)); + QueryGroupPersistenceService queryGroupPersistenceService1 = setup.v1(); + ClusterState clusterState = setup.v2(); + ClusterState newClusterState = queryGroupPersistenceService1.saveQueryGroupInClusterState(queryGroupTwo, clusterState); + Map updatedGroups = newClusterState.getMetadata().queryGroups(); + assertEquals(2, updatedGroups.size()); + assertTrue(updatedGroups.containsKey(_ID_TWO)); + Collection values = updatedGroups.values(); + assertEqualQueryGroups(queryGroupList(), new ArrayList<>(values)); + } + + /** + * Test case to ensure the error is thrown when we try to create another QueryGroup with duplicate name + */ + public void testCreateQueryGroupDuplicateName() { + Tuple setup = preparePersistenceServiceSetup(Map.of(_ID_ONE, queryGroupOne)); + QueryGroupPersistenceService queryGroupPersistenceService1 = setup.v1(); + ClusterState clusterState = setup.v2(); + QueryGroup toCreate = builder().name(NAME_ONE) + ._id("W5iIqHyhgi4K1qIAAAAIHw==") + .mode(MONITOR_STRING) + .resourceLimits(Map.of(ResourceType.fromName(MEMORY_STRING), 0.3)) + .updatedAt(1690934400000L) + .build(); + assertThrows(RuntimeException.class, () -> queryGroupPersistenceService1.saveQueryGroupInClusterState(toCreate, clusterState)); + } + + /** + * Test case to ensure the error is thrown when we try to create another QueryGroup that will make + * the total resource limits go above 1 + */ + public void testCreateQueryGroupOverflowAllocation() { + Tuple setup = preparePersistenceServiceSetup(Map.of(_ID_TWO, queryGroupTwo)); + QueryGroup toCreate = builder().name(NAME_ONE) + ._id("W5iIqHyhgi4K1qIAAAAIHw==") + .mode(MONITOR_STRING) + .resourceLimits(Map.of(ResourceType.fromName(MEMORY_STRING), 0.41)) + .updatedAt(1690934400000L) + .build(); + + QueryGroupPersistenceService queryGroupPersistenceService1 = setup.v1(); + ClusterState clusterState = setup.v2(); + assertThrows(RuntimeException.class, () -> queryGroupPersistenceService1.saveQueryGroupInClusterState(toCreate, clusterState)); + } + + /** + * Test case to ensure the error is thrown when we already have the max allowed number of QueryGroups, but + * we want to create another one + */ + public void testCreateQueryGroupOverflowCount() { + QueryGroup toCreate = builder().name(NAME_NONE_EXISTED) + ._id("W5iIqHyhgi4K1qIAAAAIHw==") + .mode(MONITOR_STRING) + .resourceLimits(Map.of(ResourceType.fromName(MEMORY_STRING), 0.5)) + .updatedAt(1690934400000L) + .build(); + Metadata metadata = Metadata.builder().queryGroups(Map.of(_ID_ONE, queryGroupOne, _ID_TWO, queryGroupTwo)).build(); + Settings settings = Settings.builder().put(QUERY_GROUP_COUNT_SETTING_NAME, 2).build(); + ClusterSettings clusterSettings = new ClusterSettings(settings, clusterSettingsSet()); + ClusterService clusterService = new ClusterService(settings, clusterSettings, mock(ThreadPool.class)); + ClusterState clusterState = ClusterState.builder(new ClusterName("_name")).metadata(metadata).build(); + QueryGroupPersistenceService queryGroupPersistenceService1 = new QueryGroupPersistenceService( + clusterService, + settings, + clusterSettings + ); + assertThrows(RuntimeException.class, () -> queryGroupPersistenceService1.saveQueryGroupInClusterState(toCreate, clusterState)); + } + + /** + * Tests the invalid value of {@code node.query_group.max_count} + */ + public void testInvalidMaxQueryGroupCount() { + Settings settings = Settings.builder().put(QUERY_GROUP_COUNT_SETTING_NAME, 2).build(); + ClusterSettings clusterSettings = new ClusterSettings(settings, clusterSettingsSet()); + ClusterService clusterService = new ClusterService(settings, clusterSettings, mock(ThreadPool.class)); + QueryGroupPersistenceService queryGroupPersistenceService = new QueryGroupPersistenceService( + clusterService, + settings, + clusterSettings + ); + assertThrows(IllegalArgumentException.class, () -> queryGroupPersistenceService.setMaxQueryGroupCount(-1)); + } + + /** + * Tests the valid value of {@code node.query_group.max_count} + */ + public void testValidMaxSandboxCountSetting() { + Settings settings = Settings.builder().put(QUERY_GROUP_COUNT_SETTING_NAME, 100).build(); + ClusterService clusterService = new ClusterService(settings, clusterSettings(), mock(ThreadPool.class)); + QueryGroupPersistenceService queryGroupPersistenceService = new QueryGroupPersistenceService( + clusterService, + settings, + clusterSettings() + ); + queryGroupPersistenceService.setMaxQueryGroupCount(50); + assertEquals(50, queryGroupPersistenceService.getMaxQueryGroupCount()); + } + + /** + * Tests PersistInClusterStateMetadata function + */ + public void testPersistInClusterStateMetadata() { + ClusterService clusterService = mock(ClusterService.class); + @SuppressWarnings("unchecked") + ActionListener listener = mock(ActionListener.class); + QueryGroupPersistenceService queryGroupPersistenceService = new QueryGroupPersistenceService( + clusterService, + QueryGroupTestUtils.settings(), + clusterSettings() + ); + queryGroupPersistenceService.persistInClusterStateMetadata(queryGroupOne, listener); + verify(clusterService).submitStateUpdateTask(eq(SOURCE), any()); + } + + /** + * Tests PersistInClusterStateMetadata function with inner functions + */ + public void testPersistInClusterStateMetadataInner() { + ClusterService clusterService = mock(ClusterService.class); + @SuppressWarnings("unchecked") + ActionListener listener = mock(ActionListener.class); + QueryGroupPersistenceService queryGroupPersistenceService = new QueryGroupPersistenceService( + clusterService, + QueryGroupTestUtils.settings(), + clusterSettings() + ); + ArgumentCaptor captor = ArgumentCaptor.forClass(ClusterStateUpdateTask.class); + queryGroupPersistenceService.persistInClusterStateMetadata(queryGroupOne, listener); + verify(clusterService, times(1)).submitStateUpdateTask(eq(SOURCE), captor.capture()); + ClusterStateUpdateTask capturedTask = captor.getValue(); + assertEquals(queryGroupPersistenceService.createQueryGroupThrottlingKey, capturedTask.getClusterManagerThrottlingKey()); + + doAnswer(invocation -> { + ClusterStateUpdateTask task = invocation.getArgument(1); + task.clusterStateProcessed(SOURCE, mock(ClusterState.class), mock(ClusterState.class)); + return null; + }).when(clusterService).submitStateUpdateTask(anyString(), any()); + queryGroupPersistenceService.persistInClusterStateMetadata(queryGroupOne, listener); + verify(listener).onResponse(any(CreateQueryGroupResponse.class)); + } + + /** + * Tests PersistInClusterStateMetadata function with failure + */ + public void testPersistInClusterStateMetadataFailure() { + ClusterService clusterService = mock(ClusterService.class); + @SuppressWarnings("unchecked") + ActionListener listener = mock(ActionListener.class); + QueryGroupPersistenceService queryGroupPersistenceService = new QueryGroupPersistenceService( + clusterService, + QueryGroupTestUtils.settings(), + clusterSettings() + ); + doAnswer(invocation -> { + ClusterStateUpdateTask task = invocation.getArgument(1); + Exception exception = new RuntimeException("Test Exception"); + task.onFailure(SOURCE, exception); + return null; + }).when(clusterService).submitStateUpdateTask(anyString(), any()); + queryGroupPersistenceService.persistInClusterStateMetadata(queryGroupOne, listener); + verify(listener).onFailure(any(RuntimeException.class)); + } +} diff --git a/plugins/workload-management/src/yamlRestTest/java/org/opensearch/plugin/wlm/WorkloadManagementClientYamlTestSuiteIT.java b/plugins/workload-management/src/yamlRestTest/java/org/opensearch/plugin/wlm/WorkloadManagementClientYamlTestSuiteIT.java new file mode 100644 index 0000000000000..9ec4a36ff6a5b --- /dev/null +++ b/plugins/workload-management/src/yamlRestTest/java/org/opensearch/plugin/wlm/WorkloadManagementClientYamlTestSuiteIT.java @@ -0,0 +1,52 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +/* + * Licensed to Elasticsearch under one or more contributor + * license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright + * ownership. Elasticsearch licenses this file to you under + * the Apache License, Version 2.0 (the "License"); you may + * not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +/* + * Modifications Copyright OpenSearch Contributors. See + * GitHub history for details. + */ + +package org.opensearch.plugin.wlm; + +import com.carrotsearch.randomizedtesting.annotations.Name; +import com.carrotsearch.randomizedtesting.annotations.ParametersFactory; + +import org.opensearch.test.rest.yaml.ClientYamlTestCandidate; +import org.opensearch.test.rest.yaml.OpenSearchClientYamlSuiteTestCase; + +/** Runs yaml rest tests */ +public class WorkloadManagementClientYamlTestSuiteIT extends OpenSearchClientYamlSuiteTestCase { + + public WorkloadManagementClientYamlTestSuiteIT(@Name("yaml") ClientYamlTestCandidate testCandidate) { + super(testCandidate); + } + + @ParametersFactory + public static Iterable parameters() throws Exception { + return OpenSearchClientYamlSuiteTestCase.createParameters(); + } +} diff --git a/plugins/workload-management/src/yamlRestTest/resources/rest-api-spec/api/create_query_group_context.json b/plugins/workload-management/src/yamlRestTest/resources/rest-api-spec/api/create_query_group_context.json new file mode 100644 index 0000000000000..bb4620c01f2d6 --- /dev/null +++ b/plugins/workload-management/src/yamlRestTest/resources/rest-api-spec/api/create_query_group_context.json @@ -0,0 +1,18 @@ +{ + "create_query_group_context": { + "stability": "experimental", + "url": { + "paths": [ + { + "path": "/_wlm/query_group", + "methods": ["PUT", "POST"], + "parts": {} + } + ] + }, + "params":{}, + "body":{ + "description":"The QueryGroup schema" + } + } +} diff --git a/plugins/workload-management/src/yamlRestTest/resources/rest-api-spec/test/wlm/10_create_query_group.yml b/plugins/workload-management/src/yamlRestTest/resources/rest-api-spec/test/wlm/10_create_query_group.yml new file mode 100644 index 0000000000000..ae82a8146e9cd --- /dev/null +++ b/plugins/workload-management/src/yamlRestTest/resources/rest-api-spec/test/wlm/10_create_query_group.yml @@ -0,0 +1,90 @@ +"test create QueryGroup API": + - skip: + version: " - 2.16.99" + reason: "QueryGroup WorkloadManagement feature was added in 2.17" + + - do: + create_query_group_context: + body: + { + "name": "analytics", + "resiliency_mode": "monitor", + "resource_limits": { + "cpu": 0.4, + "memory": 0.2 + } + } + + - match: { name: "analytics" } + - match: { resiliency_mode: "monitor" } + - match: { resource_limits.cpu: 0.4 } + - match: { resource_limits.memory: 0.2 } + + - do: + catch: /illegal_argument_exception/ + create_query_group_context: + body: + { + "name": "analytics", + "resiliency_mode": "monitor", + "resource_limits": { + "cpu": 0.4, + "memory": 0.2 + } + } + + - do: + catch: /illegal_argument_exception/ + create_query_group_context: + body: + { + "name": "analytics2", + "resiliency_mode": "monitor", + "resource_limits": { + "cpu": 0.61, + "memory": 0.2 + } + } + + - do: + catch: /illegal_argument_exception/ + create_query_group_context: + body: + { + "name": "analytics2", + "resiliency_mode": "monitor", + "resource_limits": { + "cpu": -0.1, + "memory": 0.2 + } + } + + - do: + catch: /illegal_argument_exception/ + create_query_group_context: + body: + { + "name": "", + "resiliency_mode": "monitor", + "resource_limits": { + "cpu": 0.1, + "memory": 0.2 + } + } + + - do: + create_query_group_context: + body: + { + "name": "analytics2", + "resiliency_mode": "monitor", + "resource_limits": { + "cpu": 0.35, + "memory": 0.25 + } + } + + - match: { name: "analytics2" } + - match: { resiliency_mode: "monitor" } + - match: { resource_limits.cpu: 0.35 } + - match: { resource_limits.memory: 0.25 } diff --git a/server/src/main/java/org/opensearch/cluster/metadata/Metadata.java b/server/src/main/java/org/opensearch/cluster/metadata/Metadata.java index 09bef2ddf9ee6..4da6c68b40733 100644 --- a/server/src/main/java/org/opensearch/cluster/metadata/Metadata.java +++ b/server/src/main/java/org/opensearch/cluster/metadata/Metadata.java @@ -853,6 +853,12 @@ public Map views() { return Optional.ofNullable((ViewMetadata) this.custom(ViewMetadata.TYPE)).map(ViewMetadata::views).orElse(Collections.emptyMap()); } + public Map queryGroups() { + return Optional.ofNullable((QueryGroupMetadata) this.custom(QueryGroupMetadata.TYPE)) + .map(QueryGroupMetadata::queryGroups) + .orElse(Collections.emptyMap()); + } + public DecommissionAttributeMetadata decommissionAttributeMetadata() { return custom(DecommissionAttributeMetadata.TYPE); } diff --git a/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java b/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java index beaab198073df..6ab11b1d6f150 100644 --- a/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java +++ b/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java @@ -29,37 +29,42 @@ * Class to define the QueryGroup schema * { * "_id": "fafjafjkaf9ag8a9ga9g7ag0aagaga", - * "resourceLimits": { - * "jvm": 0.4 + * "resource_limits": { + * "memory": 0.4 * }, * "resiliency_mode": "enforced", * "name": "analytics", - * "updatedAt": 4513232415 + * "updated_at": 4513232415 * } */ @ExperimentalApi public class QueryGroup extends AbstractDiffable implements ToXContentObject { + public static final String _ID_STRING = "_id"; + public static final String NAME_STRING = "name"; + public static final String RESILIENCY_MODE_STRING = "resiliency_mode"; + public static final String UPDATED_AT_STRING = "updated_at"; + public static final String RESOURCE_LIMITS_STRING = "resource_limits"; private static final int MAX_CHARS_ALLOWED_IN_NAME = 50; private final String name; private final String _id; private final ResiliencyMode resiliencyMode; // It is an epoch in millis private final long updatedAtInMillis; - private final Map resourceLimits; + private final Map resourceLimits; - public QueryGroup(String name, ResiliencyMode resiliencyMode, Map resourceLimits) { + public QueryGroup(String name, ResiliencyMode resiliencyMode, Map resourceLimits) { this(name, UUIDs.randomBase64UUID(), resiliencyMode, resourceLimits, Instant.now().getMillis()); } - public QueryGroup(String name, String _id, ResiliencyMode resiliencyMode, Map resourceLimits, long updatedAt) { + public QueryGroup(String name, String _id, ResiliencyMode resiliencyMode, Map resourceLimits, long updatedAt) { Objects.requireNonNull(name, "QueryGroup.name can't be null"); Objects.requireNonNull(resourceLimits, "QueryGroup.resourceLimits can't be null"); Objects.requireNonNull(resiliencyMode, "QueryGroup.resiliencyMode can't be null"); Objects.requireNonNull(_id, "QueryGroup._id can't be null"); - if (name.length() > MAX_CHARS_ALLOWED_IN_NAME) { - throw new IllegalArgumentException("QueryGroup.name shouldn't be more than 50 chars long"); + if (name.length() > MAX_CHARS_ALLOWED_IN_NAME || name.isEmpty()) { + throw new IllegalArgumentException("QueryGroup.name shouldn't be empty or more than 50 chars long"); } if (resourceLimits.isEmpty()) { @@ -92,7 +97,7 @@ public QueryGroup(StreamInput in) throws IOException { in.readString(), in.readString(), ResiliencyMode.fromName(in.readString()), - in.readMap((i) -> ResourceType.fromName(i.readString()), StreamInput::readGenericValue), + in.readMap((i) -> ResourceType.fromName(i.readString()), StreamInput::readDouble), in.readLong() ); } @@ -102,18 +107,18 @@ public void writeTo(StreamOutput out) throws IOException { out.writeString(name); out.writeString(_id); out.writeString(resiliencyMode.getName()); - out.writeMap(resourceLimits, ResourceType::writeTo, StreamOutput::writeGenericValue); + out.writeMap(resourceLimits, ResourceType::writeTo, StreamOutput::writeDouble); out.writeLong(updatedAtInMillis); } - private void validateResourceLimits(Map resourceLimits) { - for (Map.Entry resource : resourceLimits.entrySet()) { - Double threshold = (Double) resource.getValue(); + private void validateResourceLimits(Map resourceLimits) { + for (Map.Entry resource : resourceLimits.entrySet()) { + Double threshold = resource.getValue(); Objects.requireNonNull(resource.getKey(), "resourceName can't be null"); Objects.requireNonNull(threshold, "resource limit threshold for" + resource.getKey().getName() + " : can't be null"); - if (Double.compare(threshold, 1.0) > 0) { - throw new IllegalArgumentException("resource value should be less than 1.0"); + if (Double.compare(threshold, 0.0) <= 0 || Double.compare(threshold, 1.0) > 0) { + throw new IllegalArgumentException("resource value should be greater than 0 and less or equal to 1.0"); } } } @@ -121,12 +126,12 @@ private void validateResourceLimits(Map resourceLimits) { @Override public XContentBuilder toXContent(final XContentBuilder builder, final Params params) throws IOException { builder.startObject(); - builder.field("_id", _id); - builder.field("name", name); - builder.field("resiliency_mode", resiliencyMode.getName()); - builder.field("updatedAt", updatedAtInMillis); + builder.field(_ID_STRING, _id); + builder.field(NAME_STRING, name); + builder.field(RESILIENCY_MODE_STRING, resiliencyMode.getName()); + builder.field(UPDATED_AT_STRING, updatedAtInMillis); // write resource limits - builder.startObject("resourceLimits"); + builder.startObject(RESOURCE_LIMITS_STRING); for (ResourceType resourceType : ResourceType.values()) { if (resourceLimits.containsKey(resourceType)) { builder.field(resourceType.getName(), resourceLimits.get(resourceType)); @@ -139,56 +144,7 @@ public XContentBuilder toXContent(final XContentBuilder builder, final Params pa } public static QueryGroup fromXContent(final XContentParser parser) throws IOException { - if (parser.currentToken() == null) { // fresh parser? move to the first token - parser.nextToken(); - } - - Builder builder = builder(); - - XContentParser.Token token = parser.currentToken(); - - if (token != XContentParser.Token.START_OBJECT) { - throw new IllegalArgumentException("Expected START_OBJECT token but found [" + parser.currentName() + "]"); - } - - String fieldName = ""; - // Map to hold resources - final Map resourceLimits = new HashMap<>(); - while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) { - if (token == XContentParser.Token.FIELD_NAME) { - fieldName = parser.currentName(); - } else if (token.isValue()) { - if (fieldName.equals("_id")) { - builder._id(parser.text()); - } else if (fieldName.equals("name")) { - builder.name(parser.text()); - } else if (fieldName.equals("resiliency_mode")) { - builder.mode(parser.text()); - } else if (fieldName.equals("updatedAt")) { - builder.updatedAt(parser.longValue()); - } else { - throw new IllegalArgumentException(fieldName + " is not a valid field in QueryGroup"); - } - } else if (token == XContentParser.Token.START_OBJECT) { - - if (!fieldName.equals("resourceLimits")) { - throw new IllegalArgumentException( - "QueryGroup.resourceLimits is an object and expected token was { " + " but found " + token - ); - } - - while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) { - if (token == XContentParser.Token.FIELD_NAME) { - fieldName = parser.currentName(); - } else { - resourceLimits.put(ResourceType.fromName(fieldName), parser.doubleValue()); - } - } - - } - } - builder.resourceLimits(resourceLimits); - return builder.build(); + return Builder.fromXContent(parser).build(); } public static Diff readDiff(final StreamInput in) throws IOException { @@ -201,6 +157,7 @@ public boolean equals(Object o) { if (o == null || getClass() != o.getClass()) return false; QueryGroup that = (QueryGroup) o; return Objects.equals(name, that.name) + && Objects.equals(resiliencyMode, that.resiliencyMode) && Objects.equals(resourceLimits, that.resourceLimits) && Objects.equals(_id, that._id) && updatedAtInMillis == that.updatedAtInMillis; @@ -219,7 +176,7 @@ public ResiliencyMode getResiliencyMode() { return resiliencyMode; } - public Map getResourceLimits() { + public Map getResourceLimits() { return resourceLimits; } @@ -268,7 +225,6 @@ public static ResiliencyMode fromName(String s) { } throw new IllegalArgumentException("Invalid value for QueryGroupMode: " + s); } - } /** @@ -280,10 +236,62 @@ public static class Builder { private String _id; private ResiliencyMode resiliencyMode; private long updatedAt; - private Map resourceLimits; + private Map resourceLimits; private Builder() {} + public static Builder fromXContent(XContentParser parser) throws IOException { + if (parser.currentToken() == null) { // fresh parser? move to the first token + parser.nextToken(); + } + + Builder builder = builder(); + + XContentParser.Token token = parser.currentToken(); + + if (token != XContentParser.Token.START_OBJECT) { + throw new IllegalArgumentException("Expected START_OBJECT token but found [" + parser.currentName() + "]"); + } + + String fieldName = ""; + // Map to hold resources + final Map resourceLimits = new HashMap<>(); + while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) { + if (token == XContentParser.Token.FIELD_NAME) { + fieldName = parser.currentName(); + } else if (token.isValue()) { + if (fieldName.equals(_ID_STRING)) { + builder._id(parser.text()); + } else if (fieldName.equals(NAME_STRING)) { + builder.name(parser.text()); + } else if (fieldName.equals(RESILIENCY_MODE_STRING)) { + builder.mode(parser.text()); + } else if (fieldName.equals(UPDATED_AT_STRING)) { + builder.updatedAt(parser.longValue()); + } else { + throw new IllegalArgumentException(fieldName + " is not a valid field in QueryGroup"); + } + } else if (token == XContentParser.Token.START_OBJECT) { + + if (!fieldName.equals(RESOURCE_LIMITS_STRING)) { + throw new IllegalArgumentException( + "QueryGroup.resourceLimits is an object and expected token was { " + " but found " + token + ); + } + + while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) { + if (token == XContentParser.Token.FIELD_NAME) { + fieldName = parser.currentName(); + } else { + resourceLimits.put(ResourceType.fromName(fieldName), parser.doubleValue()); + } + } + + } + } + return builder.resourceLimits(resourceLimits); + } + public Builder name(String name) { this.name = name; return this; @@ -304,7 +312,7 @@ public Builder updatedAt(long updatedAt) { return this; } - public Builder resourceLimits(Map resourceLimits) { + public Builder resourceLimits(Map resourceLimits) { this.resourceLimits = resourceLimits; return this; } @@ -312,6 +320,5 @@ public Builder resourceLimits(Map resourceLimits) { public QueryGroup build() { return new QueryGroup(name, _id, resiliencyMode, resourceLimits, updatedAt); } - } } diff --git a/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupMetadataTests.java b/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupMetadataTests.java index d70a9ce5e10cd..06734b8e0bac2 100644 --- a/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupMetadataTests.java +++ b/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupMetadataTests.java @@ -44,7 +44,7 @@ public void testToXContent() throws IOException { queryGroupMetadata.toXContent(builder, null); builder.endObject(); assertEquals( - "{\"ajakgakg983r92_4242\":{\"_id\":\"ajakgakg983r92_4242\",\"name\":\"test\",\"resiliency_mode\":\"enforced\",\"updatedAt\":1720047207,\"resourceLimits\":{\"memory\":0.5}}}", + "{\"ajakgakg983r92_4242\":{\"_id\":\"ajakgakg983r92_4242\",\"name\":\"test\",\"resiliency_mode\":\"enforced\",\"updated_at\":1720047207,\"resource_limits\":{\"memory\":0.5}}}", builder.toString() ); } diff --git a/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupTests.java b/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupTests.java index c564f0778e6f0..884b364fb26b8 100644 --- a/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupTests.java +++ b/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupTests.java @@ -34,7 +34,7 @@ public class QueryGroupTests extends AbstractSerializingTestCase { static QueryGroup createRandomQueryGroup(String _id) { String name = randomAlphaOfLength(10); - Map resourceLimit = new HashMap<>(); + Map resourceLimit = new HashMap<>(); resourceLimit.put(ResourceType.MEMORY, randomDoubleBetween(0.0, 0.80, false)); return new QueryGroup(name, _id, randomMode(), resourceLimit, Instant.now().getMillis()); } @@ -99,10 +99,33 @@ public void testEmptyResourceLimits() { public void testIllegalQueryGroupMode() { assertThrows( NullPointerException.class, - () -> new QueryGroup("analytics", "_id", null, Map.of(ResourceType.MEMORY, (Object) 0.4), Instant.now().getMillis()) + () -> new QueryGroup("analytics", "_id", null, Map.of(ResourceType.MEMORY, 0.4), Instant.now().getMillis()) ); } + public void testQueryGroupInitiation() { + QueryGroup queryGroup = new QueryGroup("analytics", randomMode(), Map.of(ResourceType.MEMORY, 0.4)); + assertNotNull(queryGroup.getName()); + assertNotNull(queryGroup.get_id()); + assertNotNull(queryGroup.getResourceLimits()); + assertFalse(queryGroup.getResourceLimits().isEmpty()); + assertEquals(1, queryGroup.getResourceLimits().size()); + assertTrue(allowedModes.contains(queryGroup.getResiliencyMode())); + assertTrue(queryGroup.getUpdatedAtInMillis() != 0); + } + + public void testIllegalQueryGroupName() { + assertThrows( + NullPointerException.class, + () -> new QueryGroup("a".repeat(51), "_id", null, Map.of(ResourceType.MEMORY, 0.4), Instant.now().getMillis()) + ); + assertThrows( + NullPointerException.class, + () -> new QueryGroup("", "_id", null, Map.of(ResourceType.MEMORY, 0.4), Instant.now().getMillis()) + ); + + } + public void testInvalidResourceLimitWhenInvalidSystemResourceValueIsGiven() { assertThrows( IllegalArgumentException.class, @@ -110,7 +133,7 @@ public void testInvalidResourceLimitWhenInvalidSystemResourceValueIsGiven() { "analytics", "_id", randomMode(), - Map.of(ResourceType.MEMORY, (Object) randomDoubleBetween(1.1, 1.8, false)), + Map.of(ResourceType.MEMORY, randomDoubleBetween(1.1, 1.8, false)), Instant.now().getMillis() ) ); @@ -149,9 +172,9 @@ public void testToXContent() throws IOException { assertEquals( "{\"_id\":\"" + queryGroupId - + "\",\"name\":\"TestQueryGroup\",\"resiliency_mode\":\"enforced\",\"updatedAt\":" + + "\",\"name\":\"TestQueryGroup\",\"resiliency_mode\":\"enforced\",\"updated_at\":" + currentTimeInMillis - + ",\"resourceLimits\":{\"cpu\":0.3,\"memory\":0.4}}", + + ",\"resource_limits\":{\"cpu\":0.3,\"memory\":0.4}}", builder.toString() ); } From b6c80b11fa5b286fe94943dc67c3f152ff45597f Mon Sep 17 00:00:00 2001 From: Bukhtawar Khan Date: Tue, 13 Aug 2024 15:39:16 +0530 Subject: [PATCH 23/40] Refactor remote writeable entity and store to make it more reusable (#15210) * Refactor remote writeable entity and store to make it more reusable Signed-off-by: Bukhtawar Khan --- .../InternalRemoteRoutingTableService.java | 8 ++-- ...actClusterMetadataWriteableBlobEntity.java | 43 +++++++++++++++++++ .../AbstractRemoteWritableEntityManager.java | 10 ++--- .../remote/RemoteWritableEntityManager.java | 4 +- ...ty.java => RemoteWriteableBlobEntity.java} | 27 ++---------- .../RemoteWriteableEntityBlobStore.java} | 32 +++++++------- .../RemoteClusterStateAttributesManager.java | 25 ++++++----- .../remote/RemoteGlobalMetadataManager.java | 43 +++++++++++-------- .../remote/RemoteIndexMetadataManager.java | 13 +++--- .../gateway/remote/RemoteManifestManager.java | 9 ++-- .../remote/model/RemoteClusterBlocks.java | 4 +- .../model/RemoteClusterMetadataManifest.java | 4 +- .../model/RemoteClusterStateCustoms.java | 4 +- .../model/RemoteCoordinationMetadata.java | 4 +- .../remote/model/RemoteCustomMetadata.java | 4 +- .../remote/model/RemoteDiscoveryNodes.java | 4 +- .../remote/model/RemoteGlobalMetadata.java | 4 +- .../RemoteHashesOfConsistentSettings.java | 4 +- .../remote/model/RemoteIndexMetadata.java | 4 +- .../RemotePersistentSettingsMetadata.java | 4 +- .../model/RemoteRoutingTableBlobStore.java | 20 ++++++--- .../remote/model/RemoteTemplatesMetadata.java | 4 +- .../RemoteTransientSettingsMetadata.java | 4 +- .../routingtable/RemoteIndexRoutingTable.java | 4 +- .../routingtable/RemoteRoutingTableDiff.java | 4 +- ...tractRemoteWritableEntityManagerTests.java | 8 ++-- .../RemoteClusterStateServiceTests.java | 6 +-- 27 files changed, 176 insertions(+), 128 deletions(-) create mode 100644 server/src/main/java/org/opensearch/common/remote/AbstractClusterMetadataWriteableBlobEntity.java rename server/src/main/java/org/opensearch/common/remote/{AbstractRemoteWritableBlobEntity.java => RemoteWriteableBlobEntity.java} (62%) rename server/src/main/java/org/opensearch/{gateway/remote/model/RemoteClusterStateBlobStore.java => common/remote/RemoteWriteableEntityBlobStore.java} (78%) diff --git a/server/src/main/java/org/opensearch/cluster/routing/remote/InternalRemoteRoutingTableService.java b/server/src/main/java/org/opensearch/cluster/routing/remote/InternalRemoteRoutingTableService.java index 3c578a8c5c01f..0f1ff3138ef90 100644 --- a/server/src/main/java/org/opensearch/cluster/routing/remote/InternalRemoteRoutingTableService.java +++ b/server/src/main/java/org/opensearch/cluster/routing/remote/InternalRemoteRoutingTableService.java @@ -20,14 +20,15 @@ import org.opensearch.common.blobstore.BlobPath; import org.opensearch.common.lifecycle.AbstractLifecycleComponent; import org.opensearch.common.remote.RemoteWritableEntityStore; +import org.opensearch.common.remote.RemoteWriteableEntityBlobStore; import org.opensearch.common.settings.ClusterSettings; import org.opensearch.common.settings.Settings; import org.opensearch.common.util.io.IOUtils; import org.opensearch.core.action.ActionListener; import org.opensearch.core.compress.Compressor; import org.opensearch.gateway.remote.ClusterMetadataManifest; +import org.opensearch.gateway.remote.RemoteClusterStateUtils; import org.opensearch.gateway.remote.RemoteStateTransferException; -import org.opensearch.gateway.remote.model.RemoteClusterStateBlobStore; import org.opensearch.gateway.remote.model.RemoteRoutingTableBlobStore; import org.opensearch.gateway.remote.routingtable.RemoteIndexRoutingTable; import org.opensearch.gateway.remote.routingtable.RemoteRoutingTableDiff; @@ -262,12 +263,13 @@ protected void doStart() { clusterSettings ); - this.remoteRoutingTableDiffStore = new RemoteClusterStateBlobStore<>( + this.remoteRoutingTableDiffStore = new RemoteWriteableEntityBlobStore<>( new BlobStoreTransferService(blobStoreRepository.blobStore(), threadPool), blobStoreRepository, clusterName, threadPool, - ThreadPool.Names.REMOTE_STATE_READ + ThreadPool.Names.REMOTE_STATE_READ, + RemoteClusterStateUtils.CLUSTER_STATE_PATH_TOKEN ); } diff --git a/server/src/main/java/org/opensearch/common/remote/AbstractClusterMetadataWriteableBlobEntity.java b/server/src/main/java/org/opensearch/common/remote/AbstractClusterMetadataWriteableBlobEntity.java new file mode 100644 index 0000000000000..b9492da04680d --- /dev/null +++ b/server/src/main/java/org/opensearch/common/remote/AbstractClusterMetadataWriteableBlobEntity.java @@ -0,0 +1,43 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.common.remote; + +import org.opensearch.core.compress.Compressor; +import org.opensearch.core.xcontent.NamedXContentRegistry; +import org.opensearch.gateway.remote.ClusterMetadataManifest.UploadedMetadata; + +/** + * An extension of {@link RemoteWriteableEntity} class which caters to the use case of writing to and reading from a blob storage + * + * @param The class type which can be uploaded to or downloaded from a blob storage. + */ +public abstract class AbstractClusterMetadataWriteableBlobEntity extends RemoteWriteableBlobEntity { + + protected final NamedXContentRegistry namedXContentRegistry; + + public AbstractClusterMetadataWriteableBlobEntity( + final String clusterUUID, + final Compressor compressor, + final NamedXContentRegistry namedXContentRegistry + ) { + super(clusterUUID, compressor); + this.namedXContentRegistry = namedXContentRegistry; + } + + public AbstractClusterMetadataWriteableBlobEntity(final String clusterUUID, final Compressor compressor) { + super(clusterUUID, compressor); + this.namedXContentRegistry = null; + } + + public abstract UploadedMetadata getUploadedMetadata(); + + public NamedXContentRegistry getNamedXContentRegistry() { + return namedXContentRegistry; + } +} diff --git a/server/src/main/java/org/opensearch/common/remote/AbstractRemoteWritableEntityManager.java b/server/src/main/java/org/opensearch/common/remote/AbstractRemoteWritableEntityManager.java index dc301635c4a80..8e2de1580a49f 100644 --- a/server/src/main/java/org/opensearch/common/remote/AbstractRemoteWritableEntityManager.java +++ b/server/src/main/java/org/opensearch/common/remote/AbstractRemoteWritableEntityManager.java @@ -31,7 +31,7 @@ public abstract class AbstractRemoteWritableEntityManager implements RemoteWrita * @return the remote writable entity store for the given entity * @throws IllegalArgumentException if the entity type is unknown */ - protected RemoteWritableEntityStore getStore(AbstractRemoteWritableBlobEntity entity) { + protected RemoteWritableEntityStore getStore(AbstractClusterMetadataWriteableBlobEntity entity) { RemoteWritableEntityStore remoteStore = remoteWritableEntityStores.get(entity.getType()); if (remoteStore == null) { throw new IllegalArgumentException("Unknown entity type [" + entity.getType() + "]"); @@ -49,7 +49,7 @@ protected RemoteWritableEntityStore getStore(AbstractRemoteWritableBlobEntity en */ protected abstract ActionListener getWrappedWriteListener( String component, - AbstractRemoteWritableBlobEntity remoteEntity, + AbstractClusterMetadataWriteableBlobEntity remoteEntity, ActionListener listener ); @@ -64,21 +64,21 @@ protected abstract ActionListener getWrappedWriteListener( */ protected abstract ActionListener getWrappedReadListener( String component, - AbstractRemoteWritableBlobEntity remoteEntity, + AbstractClusterMetadataWriteableBlobEntity remoteEntity, ActionListener listener ); @Override public void writeAsync( String component, - AbstractRemoteWritableBlobEntity entity, + AbstractClusterMetadataWriteableBlobEntity entity, ActionListener listener ) { getStore(entity).writeAsync(entity, getWrappedWriteListener(component, entity, listener)); } @Override - public void readAsync(String component, AbstractRemoteWritableBlobEntity entity, ActionListener listener) { + public void readAsync(String component, AbstractClusterMetadataWriteableBlobEntity entity, ActionListener listener) { getStore(entity).readAsync(entity, getWrappedReadListener(component, entity, listener)); } } diff --git a/server/src/main/java/org/opensearch/common/remote/RemoteWritableEntityManager.java b/server/src/main/java/org/opensearch/common/remote/RemoteWritableEntityManager.java index 7693d1b5284bd..c27598e368e4d 100644 --- a/server/src/main/java/org/opensearch/common/remote/RemoteWritableEntityManager.java +++ b/server/src/main/java/org/opensearch/common/remote/RemoteWritableEntityManager.java @@ -29,7 +29,7 @@ public interface RemoteWritableEntityManager { * {@link ActionListener#onFailure(Exception)} method is called with * an exception if the read operation fails. */ - void readAsync(String component, AbstractRemoteWritableBlobEntity entity, ActionListener listener); + void readAsync(String component, AbstractClusterMetadataWriteableBlobEntity entity, ActionListener listener); /** * Performs an asynchronous write operation for the specified component and entity. @@ -43,5 +43,5 @@ public interface RemoteWritableEntityManager { * {@link ActionListener#onFailure(Exception)} method is called with * an exception if the write operation fails. */ - void writeAsync(String component, AbstractRemoteWritableBlobEntity entity, ActionListener listener); + void writeAsync(String component, AbstractClusterMetadataWriteableBlobEntity entity, ActionListener listener); } diff --git a/server/src/main/java/org/opensearch/common/remote/AbstractRemoteWritableBlobEntity.java b/server/src/main/java/org/opensearch/common/remote/RemoteWriteableBlobEntity.java similarity index 62% rename from server/src/main/java/org/opensearch/common/remote/AbstractRemoteWritableBlobEntity.java rename to server/src/main/java/org/opensearch/common/remote/RemoteWriteableBlobEntity.java index 237c077cb673c..f034ce2d1adf1 100644 --- a/server/src/main/java/org/opensearch/common/remote/AbstractRemoteWritableBlobEntity.java +++ b/server/src/main/java/org/opensearch/common/remote/RemoteWriteableBlobEntity.java @@ -10,38 +10,25 @@ import org.opensearch.common.blobstore.BlobPath; import org.opensearch.core.compress.Compressor; -import org.opensearch.core.xcontent.NamedXContentRegistry; -import org.opensearch.gateway.remote.ClusterMetadataManifest.UploadedMetadata; import static org.opensearch.gateway.remote.RemoteClusterStateUtils.PATH_DELIMITER; /** - * An extension of {@link RemoteWriteableEntity} class which caters to the use case of writing to and reading from a blob storage - * - * @param The class type which can be uploaded to or downloaded from a blob storage. + * The abstract class which represents a {@link RemoteWriteableEntity} that can be written to a store + * @param the entity to be written */ -public abstract class AbstractRemoteWritableBlobEntity implements RemoteWriteableEntity { +public abstract class RemoteWriteableBlobEntity implements RemoteWriteableEntity { protected String blobFileName; protected String blobName; private final String clusterUUID; private final Compressor compressor; - private final NamedXContentRegistry namedXContentRegistry; private String[] pathTokens; - public AbstractRemoteWritableBlobEntity( - final String clusterUUID, - final Compressor compressor, - final NamedXContentRegistry namedXContentRegistry - ) { + public RemoteWriteableBlobEntity(final String clusterUUID, final Compressor compressor) { this.clusterUUID = clusterUUID; this.compressor = compressor; - this.namedXContentRegistry = namedXContentRegistry; - } - - public AbstractRemoteWritableBlobEntity(final String clusterUUID, final Compressor compressor) { - this(clusterUUID, compressor, null); } public abstract BlobPathParameters getBlobPathParameters(); @@ -80,16 +67,10 @@ public String clusterUUID() { return clusterUUID; } - public abstract UploadedMetadata getUploadedMetadata(); - public void setFullBlobName(BlobPath blobPath) { this.blobName = blobPath.buildAsString() + blobFileName; } - public NamedXContentRegistry getNamedXContentRegistry() { - return namedXContentRegistry; - } - protected Compressor getCompressor() { return compressor; } diff --git a/server/src/main/java/org/opensearch/gateway/remote/model/RemoteClusterStateBlobStore.java b/server/src/main/java/org/opensearch/common/remote/RemoteWriteableEntityBlobStore.java similarity index 78% rename from server/src/main/java/org/opensearch/gateway/remote/model/RemoteClusterStateBlobStore.java rename to server/src/main/java/org/opensearch/common/remote/RemoteWriteableEntityBlobStore.java index cd8b8aa41ad65..baa44a7c9bde9 100644 --- a/server/src/main/java/org/opensearch/gateway/remote/model/RemoteClusterStateBlobStore.java +++ b/server/src/main/java/org/opensearch/common/remote/RemoteWriteableEntityBlobStore.java @@ -6,49 +6,48 @@ * compatible open source license. */ -package org.opensearch.gateway.remote.model; +package org.opensearch.common.remote; import org.opensearch.common.blobstore.BlobPath; import org.opensearch.common.blobstore.stream.write.WritePriority; -import org.opensearch.common.remote.AbstractRemoteWritableBlobEntity; -import org.opensearch.common.remote.RemoteWritableEntityStore; -import org.opensearch.common.remote.RemoteWriteableEntity; import org.opensearch.core.action.ActionListener; -import org.opensearch.gateway.remote.RemoteClusterStateUtils; import org.opensearch.index.translog.transfer.BlobStoreTransferService; import org.opensearch.repositories.blobstore.BlobStoreRepository; import org.opensearch.threadpool.ThreadPool; import java.io.IOException; import java.io.InputStream; +import java.nio.charset.StandardCharsets; +import java.util.Base64; import java.util.concurrent.ExecutorService; -import static org.opensearch.gateway.remote.RemoteClusterStateUtils.CLUSTER_STATE_PATH_TOKEN; - /** * Abstract class for a blob type storage * * @param The entity which can be uploaded to / downloaded from blob store * @param The concrete class implementing {@link RemoteWriteableEntity} which is used as a wrapper for T entity. */ -public class RemoteClusterStateBlobStore> implements RemoteWritableEntityStore { +public class RemoteWriteableEntityBlobStore> implements RemoteWritableEntityStore { private final BlobStoreTransferService transferService; private final BlobStoreRepository blobStoreRepository; private final String clusterName; private final ExecutorService executorService; + private final String pathToken; - public RemoteClusterStateBlobStore( + public RemoteWriteableEntityBlobStore( final BlobStoreTransferService blobStoreTransferService, final BlobStoreRepository blobStoreRepository, final String clusterName, final ThreadPool threadPool, - final String executor + final String executor, + final String pathToken ) { this.transferService = blobStoreTransferService; this.blobStoreRepository = blobStoreRepository; this.clusterName = clusterName; this.executorService = threadPool.executor(executor); + this.pathToken = pathToken; } @Override @@ -95,13 +94,10 @@ public String getClusterName() { } public BlobPath getBlobPathPrefix(String clusterUUID) { - return blobStoreRepository.basePath() - .add(RemoteClusterStateUtils.encodeString(getClusterName())) - .add(CLUSTER_STATE_PATH_TOKEN) - .add(clusterUUID); + return blobStoreRepository.basePath().add(encodeString(getClusterName())).add(pathToken).add(clusterUUID); } - public BlobPath getBlobPathForUpload(final AbstractRemoteWritableBlobEntity obj) { + public BlobPath getBlobPathForUpload(final RemoteWriteableBlobEntity obj) { BlobPath blobPath = getBlobPathPrefix(obj.clusterUUID()); for (String token : obj.getBlobPathParameters().getPathTokens()) { blobPath = blobPath.add(token); @@ -109,7 +105,7 @@ public BlobPath getBlobPathForUpload(final AbstractRemoteWritableBlobEntity o return blobPath; } - public BlobPath getBlobPathForDownload(final AbstractRemoteWritableBlobEntity obj) { + public BlobPath getBlobPathForDownload(final RemoteWriteableBlobEntity obj) { String[] pathTokens = obj.getBlobPathTokens(); BlobPath blobPath = new BlobPath(); if (pathTokens == null || pathTokens.length < 1) { @@ -122,4 +118,8 @@ public BlobPath getBlobPathForDownload(final AbstractRemoteWritableBlobEntity return blobPath; } + private static String encodeString(String content) { + return Base64.getUrlEncoder().withoutPadding().encodeToString(content.getBytes(StandardCharsets.UTF_8)); + } + } diff --git a/server/src/main/java/org/opensearch/gateway/remote/RemoteClusterStateAttributesManager.java b/server/src/main/java/org/opensearch/gateway/remote/RemoteClusterStateAttributesManager.java index 67ac8d2b9a810..877e2585cb1eb 100644 --- a/server/src/main/java/org/opensearch/gateway/remote/RemoteClusterStateAttributesManager.java +++ b/server/src/main/java/org/opensearch/gateway/remote/RemoteClusterStateAttributesManager.java @@ -11,12 +11,12 @@ import org.opensearch.cluster.ClusterState; import org.opensearch.cluster.DiffableUtils; import org.opensearch.cluster.DiffableUtils.NonDiffableValueSerializer; -import org.opensearch.common.remote.AbstractRemoteWritableBlobEntity; +import org.opensearch.common.remote.AbstractClusterMetadataWriteableBlobEntity; import org.opensearch.common.remote.AbstractRemoteWritableEntityManager; +import org.opensearch.common.remote.RemoteWriteableEntityBlobStore; import org.opensearch.core.action.ActionListener; import org.opensearch.core.common.io.stream.NamedWriteableRegistry; import org.opensearch.gateway.remote.model.RemoteClusterBlocks; -import org.opensearch.gateway.remote.model.RemoteClusterStateBlobStore; import org.opensearch.gateway.remote.model.RemoteClusterStateCustoms; import org.opensearch.gateway.remote.model.RemoteDiscoveryNodes; import org.opensearch.gateway.remote.model.RemoteReadResult; @@ -28,7 +28,7 @@ import java.util.Map; /** - * A Manager which provides APIs to upload and download attributes of ClusterState to the {@link RemoteClusterStateBlobStore} + * A Manager which provides APIs to upload and download attributes of ClusterState to the {@link RemoteWriteableEntityBlobStore} * * @opensearch.internal */ @@ -47,32 +47,35 @@ public class RemoteClusterStateAttributesManager extends AbstractRemoteWritableE ) { this.remoteWritableEntityStores.put( RemoteDiscoveryNodes.DISCOVERY_NODES, - new RemoteClusterStateBlobStore<>( + new RemoteWriteableEntityBlobStore<>( blobStoreTransferService, blobStoreRepository, clusterName, threadpool, - ThreadPool.Names.REMOTE_STATE_READ + ThreadPool.Names.REMOTE_STATE_READ, + RemoteClusterStateUtils.CLUSTER_STATE_PATH_TOKEN ) ); this.remoteWritableEntityStores.put( RemoteClusterBlocks.CLUSTER_BLOCKS, - new RemoteClusterStateBlobStore<>( + new RemoteWriteableEntityBlobStore<>( blobStoreTransferService, blobStoreRepository, clusterName, threadpool, - ThreadPool.Names.REMOTE_STATE_READ + ThreadPool.Names.REMOTE_STATE_READ, + RemoteClusterStateUtils.CLUSTER_STATE_PATH_TOKEN ) ); this.remoteWritableEntityStores.put( RemoteClusterStateCustoms.CLUSTER_STATE_CUSTOM, - new RemoteClusterStateBlobStore<>( + new RemoteWriteableEntityBlobStore<>( blobStoreTransferService, blobStoreRepository, clusterName, threadpool, - ThreadPool.Names.REMOTE_STATE_READ + ThreadPool.Names.REMOTE_STATE_READ, + RemoteClusterStateUtils.CLUSTER_STATE_PATH_TOKEN ) ); } @@ -80,7 +83,7 @@ public class RemoteClusterStateAttributesManager extends AbstractRemoteWritableE @Override protected ActionListener getWrappedWriteListener( String component, - AbstractRemoteWritableBlobEntity remoteEntity, + AbstractClusterMetadataWriteableBlobEntity remoteEntity, ActionListener listener ) { return ActionListener.wrap( @@ -92,7 +95,7 @@ protected ActionListener getWrappedWriteListener( @Override protected ActionListener getWrappedReadListener( String component, - AbstractRemoteWritableBlobEntity remoteEntity, + AbstractClusterMetadataWriteableBlobEntity remoteEntity, ActionListener listener ) { return ActionListener.wrap( diff --git a/server/src/main/java/org/opensearch/gateway/remote/RemoteGlobalMetadataManager.java b/server/src/main/java/org/opensearch/gateway/remote/RemoteGlobalMetadataManager.java index 5a6f4b7e9f1f1..763a8e3ff4951 100644 --- a/server/src/main/java/org/opensearch/gateway/remote/RemoteGlobalMetadataManager.java +++ b/server/src/main/java/org/opensearch/gateway/remote/RemoteGlobalMetadataManager.java @@ -16,8 +16,9 @@ import org.opensearch.cluster.metadata.Metadata.Custom; import org.opensearch.cluster.metadata.Metadata.XContentContext; import org.opensearch.cluster.metadata.TemplatesMetadata; -import org.opensearch.common.remote.AbstractRemoteWritableBlobEntity; +import org.opensearch.common.remote.AbstractClusterMetadataWriteableBlobEntity; import org.opensearch.common.remote.AbstractRemoteWritableEntityManager; +import org.opensearch.common.remote.RemoteWriteableEntityBlobStore; import org.opensearch.common.settings.ClusterSettings; import org.opensearch.common.settings.Setting; import org.opensearch.common.settings.Settings; @@ -26,7 +27,6 @@ import org.opensearch.core.common.io.stream.NamedWriteableRegistry; import org.opensearch.core.compress.Compressor; import org.opensearch.core.xcontent.NamedXContentRegistry; -import org.opensearch.gateway.remote.model.RemoteClusterStateBlobStore; import org.opensearch.gateway.remote.model.RemoteCoordinationMetadata; import org.opensearch.gateway.remote.model.RemoteCustomMetadata; import org.opensearch.gateway.remote.model.RemoteGlobalMetadata; @@ -85,72 +85,79 @@ public class RemoteGlobalMetadataManager extends AbstractRemoteWritableEntityMan this.namedWriteableRegistry = namedWriteableRegistry; this.remoteWritableEntityStores.put( RemoteGlobalMetadata.GLOBAL_METADATA, - new RemoteClusterStateBlobStore<>( + new RemoteWriteableEntityBlobStore<>( blobStoreTransferService, blobStoreRepository, clusterName, threadpool, - ThreadPool.Names.REMOTE_STATE_READ + ThreadPool.Names.REMOTE_STATE_READ, + RemoteClusterStateUtils.CLUSTER_STATE_PATH_TOKEN ) ); this.remoteWritableEntityStores.put( RemoteCoordinationMetadata.COORDINATION_METADATA, - new RemoteClusterStateBlobStore<>( + new RemoteWriteableEntityBlobStore<>( blobStoreTransferService, blobStoreRepository, clusterName, threadpool, - ThreadPool.Names.REMOTE_STATE_READ + ThreadPool.Names.REMOTE_STATE_READ, + RemoteClusterStateUtils.CLUSTER_STATE_PATH_TOKEN ) ); this.remoteWritableEntityStores.put( RemotePersistentSettingsMetadata.SETTING_METADATA, - new RemoteClusterStateBlobStore<>( + new RemoteWriteableEntityBlobStore<>( blobStoreTransferService, blobStoreRepository, clusterName, threadpool, - ThreadPool.Names.REMOTE_STATE_READ + ThreadPool.Names.REMOTE_STATE_READ, + RemoteClusterStateUtils.CLUSTER_STATE_PATH_TOKEN ) ); this.remoteWritableEntityStores.put( RemoteTransientSettingsMetadata.TRANSIENT_SETTING_METADATA, - new RemoteClusterStateBlobStore<>( + new RemoteWriteableEntityBlobStore<>( blobStoreTransferService, blobStoreRepository, clusterName, threadpool, - ThreadPool.Names.REMOTE_STATE_READ + ThreadPool.Names.REMOTE_STATE_READ, + RemoteClusterStateUtils.CLUSTER_STATE_PATH_TOKEN ) ); this.remoteWritableEntityStores.put( RemoteHashesOfConsistentSettings.HASHES_OF_CONSISTENT_SETTINGS, - new RemoteClusterStateBlobStore<>( + new RemoteWriteableEntityBlobStore<>( blobStoreTransferService, blobStoreRepository, clusterName, threadpool, - ThreadPool.Names.REMOTE_STATE_READ + ThreadPool.Names.REMOTE_STATE_READ, + RemoteClusterStateUtils.CLUSTER_STATE_PATH_TOKEN ) ); this.remoteWritableEntityStores.put( RemoteTemplatesMetadata.TEMPLATES_METADATA, - new RemoteClusterStateBlobStore<>( + new RemoteWriteableEntityBlobStore<>( blobStoreTransferService, blobStoreRepository, clusterName, threadpool, - ThreadPool.Names.REMOTE_STATE_READ + ThreadPool.Names.REMOTE_STATE_READ, + RemoteClusterStateUtils.CLUSTER_STATE_PATH_TOKEN ) ); this.remoteWritableEntityStores.put( RemoteCustomMetadata.CUSTOM_METADATA, - new RemoteClusterStateBlobStore<>( + new RemoteWriteableEntityBlobStore<>( blobStoreTransferService, blobStoreRepository, clusterName, threadpool, - ThreadPool.Names.REMOTE_STATE_READ + ThreadPool.Names.REMOTE_STATE_READ, + RemoteClusterStateUtils.CLUSTER_STATE_PATH_TOKEN ) ); clusterSettings.addSettingsUpdateConsumer(GLOBAL_METADATA_UPLOAD_TIMEOUT_SETTING, this::setGlobalMetadataUploadTimeout); @@ -159,7 +166,7 @@ public class RemoteGlobalMetadataManager extends AbstractRemoteWritableEntityMan @Override protected ActionListener getWrappedWriteListener( String component, - AbstractRemoteWritableBlobEntity remoteEntity, + AbstractClusterMetadataWriteableBlobEntity remoteEntity, ActionListener listener ) { return ActionListener.wrap( @@ -171,7 +178,7 @@ protected ActionListener getWrappedWriteListener( @Override protected ActionListener getWrappedReadListener( String component, - AbstractRemoteWritableBlobEntity remoteEntity, + AbstractClusterMetadataWriteableBlobEntity remoteEntity, ActionListener listener ) { return ActionListener.wrap( diff --git a/server/src/main/java/org/opensearch/gateway/remote/RemoteIndexMetadataManager.java b/server/src/main/java/org/opensearch/gateway/remote/RemoteIndexMetadataManager.java index c30721c8f625c..d1f08a7c2a33d 100644 --- a/server/src/main/java/org/opensearch/gateway/remote/RemoteIndexMetadataManager.java +++ b/server/src/main/java/org/opensearch/gateway/remote/RemoteIndexMetadataManager.java @@ -9,15 +9,15 @@ package org.opensearch.gateway.remote; import org.opensearch.cluster.metadata.IndexMetadata; -import org.opensearch.common.remote.AbstractRemoteWritableBlobEntity; +import org.opensearch.common.remote.AbstractClusterMetadataWriteableBlobEntity; import org.opensearch.common.remote.AbstractRemoteWritableEntityManager; +import org.opensearch.common.remote.RemoteWriteableEntityBlobStore; import org.opensearch.common.settings.ClusterSettings; import org.opensearch.common.settings.Setting; import org.opensearch.common.unit.TimeValue; import org.opensearch.core.action.ActionListener; import org.opensearch.core.compress.Compressor; import org.opensearch.core.xcontent.NamedXContentRegistry; -import org.opensearch.gateway.remote.model.RemoteClusterStateBlobStore; import org.opensearch.gateway.remote.model.RemoteIndexMetadata; import org.opensearch.gateway.remote.model.RemoteReadResult; import org.opensearch.index.translog.transfer.BlobStoreTransferService; @@ -58,12 +58,13 @@ public RemoteIndexMetadataManager( ) { this.remoteWritableEntityStores.put( RemoteIndexMetadata.INDEX, - new RemoteClusterStateBlobStore<>( + new RemoteWriteableEntityBlobStore<>( blobStoreTransferService, blobStoreRepository, clusterName, threadpool, - ThreadPool.Names.REMOTE_STATE_READ + ThreadPool.Names.REMOTE_STATE_READ, + RemoteClusterStateUtils.CLUSTER_STATE_PATH_TOKEN ) ); this.namedXContentRegistry = blobStoreRepository.getNamedXContentRegistry(); @@ -106,7 +107,7 @@ private void setIndexMetadataUploadTimeout(TimeValue newIndexMetadataUploadTimeo @Override protected ActionListener getWrappedWriteListener( String component, - AbstractRemoteWritableBlobEntity remoteEntity, + AbstractClusterMetadataWriteableBlobEntity remoteEntity, ActionListener listener ) { return ActionListener.wrap( @@ -118,7 +119,7 @@ protected ActionListener getWrappedWriteListener( @Override protected ActionListener getWrappedReadListener( String component, - AbstractRemoteWritableBlobEntity remoteEntity, + AbstractClusterMetadataWriteableBlobEntity remoteEntity, ActionListener listener ) { return ActionListener.wrap( diff --git a/server/src/main/java/org/opensearch/gateway/remote/RemoteManifestManager.java b/server/src/main/java/org/opensearch/gateway/remote/RemoteManifestManager.java index cb09de1a6ec44..0ccadd7dd18da 100644 --- a/server/src/main/java/org/opensearch/gateway/remote/RemoteManifestManager.java +++ b/server/src/main/java/org/opensearch/gateway/remote/RemoteManifestManager.java @@ -16,6 +16,7 @@ import org.opensearch.common.blobstore.BlobContainer; import org.opensearch.common.blobstore.BlobMetadata; import org.opensearch.common.blobstore.BlobPath; +import org.opensearch.common.remote.RemoteWriteableEntityBlobStore; import org.opensearch.common.settings.ClusterSettings; import org.opensearch.common.settings.Setting; import org.opensearch.common.unit.TimeValue; @@ -23,7 +24,6 @@ import org.opensearch.core.compress.Compressor; import org.opensearch.core.xcontent.NamedXContentRegistry; import org.opensearch.gateway.remote.model.RemoteClusterMetadataManifest; -import org.opensearch.gateway.remote.model.RemoteClusterStateBlobStore; import org.opensearch.gateway.remote.model.RemoteClusterStateManifestInfo; import org.opensearch.index.remote.RemoteStoreUtils; import org.opensearch.index.translog.transfer.BlobStoreTransferService; @@ -62,7 +62,7 @@ public class RemoteManifestManager { private volatile TimeValue metadataManifestUploadTimeout; private final String nodeId; - private final RemoteClusterStateBlobStore manifestBlobStore; + private final RemoteWriteableEntityBlobStore manifestBlobStore; private final Compressor compressor; private final NamedXContentRegistry namedXContentRegistry; // todo remove blobStorerepo from here @@ -78,12 +78,13 @@ public class RemoteManifestManager { ) { this.metadataManifestUploadTimeout = clusterSettings.get(METADATA_MANIFEST_UPLOAD_TIMEOUT_SETTING); this.nodeId = nodeId; - this.manifestBlobStore = new RemoteClusterStateBlobStore<>( + this.manifestBlobStore = new RemoteWriteableEntityBlobStore<>( blobStoreTransferService, blobStoreRepository, clusterName, threadpool, - ThreadPool.Names.REMOTE_STATE_READ + ThreadPool.Names.REMOTE_STATE_READ, + RemoteClusterStateUtils.CLUSTER_STATE_PATH_TOKEN ); ; clusterSettings.addSettingsUpdateConsumer(METADATA_MANIFEST_UPLOAD_TIMEOUT_SETTING, this::setMetadataManifestUploadTimeout); diff --git a/server/src/main/java/org/opensearch/gateway/remote/model/RemoteClusterBlocks.java b/server/src/main/java/org/opensearch/gateway/remote/model/RemoteClusterBlocks.java index 9c5fbd5941640..101daaa143a66 100644 --- a/server/src/main/java/org/opensearch/gateway/remote/model/RemoteClusterBlocks.java +++ b/server/src/main/java/org/opensearch/gateway/remote/model/RemoteClusterBlocks.java @@ -10,7 +10,7 @@ import org.opensearch.cluster.block.ClusterBlocks; import org.opensearch.common.io.Streams; -import org.opensearch.common.remote.AbstractRemoteWritableBlobEntity; +import org.opensearch.common.remote.AbstractClusterMetadataWriteableBlobEntity; import org.opensearch.common.remote.BlobPathParameters; import org.opensearch.core.compress.Compressor; import org.opensearch.gateway.remote.ClusterMetadataManifest.UploadedMetadata; @@ -29,7 +29,7 @@ /** * Wrapper class for uploading/downloading {@link ClusterBlocks} to/from remote blob store */ -public class RemoteClusterBlocks extends AbstractRemoteWritableBlobEntity { +public class RemoteClusterBlocks extends AbstractClusterMetadataWriteableBlobEntity { public static final String CLUSTER_BLOCKS = "blocks"; public static final ChecksumWritableBlobStoreFormat CLUSTER_BLOCKS_FORMAT = new ChecksumWritableBlobStoreFormat<>( diff --git a/server/src/main/java/org/opensearch/gateway/remote/model/RemoteClusterMetadataManifest.java b/server/src/main/java/org/opensearch/gateway/remote/model/RemoteClusterMetadataManifest.java index acaae3173315a..5f79b690af574 100644 --- a/server/src/main/java/org/opensearch/gateway/remote/model/RemoteClusterMetadataManifest.java +++ b/server/src/main/java/org/opensearch/gateway/remote/model/RemoteClusterMetadataManifest.java @@ -9,7 +9,7 @@ package org.opensearch.gateway.remote.model; import org.opensearch.common.io.Streams; -import org.opensearch.common.remote.AbstractRemoteWritableBlobEntity; +import org.opensearch.common.remote.AbstractClusterMetadataWriteableBlobEntity; import org.opensearch.common.remote.BlobPathParameters; import org.opensearch.core.compress.Compressor; import org.opensearch.core.xcontent.NamedXContentRegistry; @@ -29,7 +29,7 @@ /** * Wrapper class for uploading/downloading {@link ClusterMetadataManifest} to/from remote blob store */ -public class RemoteClusterMetadataManifest extends AbstractRemoteWritableBlobEntity { +public class RemoteClusterMetadataManifest extends AbstractClusterMetadataWriteableBlobEntity { public static final String MANIFEST = "manifest"; public static final int SPLITTED_MANIFEST_FILE_LENGTH = 6; diff --git a/server/src/main/java/org/opensearch/gateway/remote/model/RemoteClusterStateCustoms.java b/server/src/main/java/org/opensearch/gateway/remote/model/RemoteClusterStateCustoms.java index affbc7ba66cb8..e5e44525520f4 100644 --- a/server/src/main/java/org/opensearch/gateway/remote/model/RemoteClusterStateCustoms.java +++ b/server/src/main/java/org/opensearch/gateway/remote/model/RemoteClusterStateCustoms.java @@ -11,7 +11,7 @@ import org.opensearch.cluster.ClusterState; import org.opensearch.cluster.ClusterState.Custom; import org.opensearch.common.io.Streams; -import org.opensearch.common.remote.AbstractRemoteWritableBlobEntity; +import org.opensearch.common.remote.AbstractClusterMetadataWriteableBlobEntity; import org.opensearch.common.remote.BlobPathParameters; import org.opensearch.core.common.io.stream.NamedWriteableRegistry; import org.opensearch.core.common.io.stream.StreamInput; @@ -32,7 +32,7 @@ /** * Wrapper class for uploading/downloading {@link Custom} to/from remote blob store */ -public class RemoteClusterStateCustoms extends AbstractRemoteWritableBlobEntity { +public class RemoteClusterStateCustoms extends AbstractClusterMetadataWriteableBlobEntity { public static final String CLUSTER_STATE_CUSTOM = "cluster-state-custom"; public final ChecksumWritableBlobStoreFormat clusterStateCustomsFormat; diff --git a/server/src/main/java/org/opensearch/gateway/remote/model/RemoteCoordinationMetadata.java b/server/src/main/java/org/opensearch/gateway/remote/model/RemoteCoordinationMetadata.java index a90721ab59f66..63cc96e3e02c4 100644 --- a/server/src/main/java/org/opensearch/gateway/remote/model/RemoteCoordinationMetadata.java +++ b/server/src/main/java/org/opensearch/gateway/remote/model/RemoteCoordinationMetadata.java @@ -10,7 +10,7 @@ import org.opensearch.cluster.coordination.CoordinationMetadata; import org.opensearch.common.io.Streams; -import org.opensearch.common.remote.AbstractRemoteWritableBlobEntity; +import org.opensearch.common.remote.AbstractClusterMetadataWriteableBlobEntity; import org.opensearch.common.remote.BlobPathParameters; import org.opensearch.core.compress.Compressor; import org.opensearch.core.xcontent.NamedXContentRegistry; @@ -31,7 +31,7 @@ /** * Wrapper class for uploading/downloading {@link CoordinationMetadata} to/from remote blob store */ -public class RemoteCoordinationMetadata extends AbstractRemoteWritableBlobEntity { +public class RemoteCoordinationMetadata extends AbstractClusterMetadataWriteableBlobEntity { public static final String COORDINATION_METADATA = "coordination"; public static final ChecksumBlobStoreFormat COORDINATION_METADATA_FORMAT = new ChecksumBlobStoreFormat<>( diff --git a/server/src/main/java/org/opensearch/gateway/remote/model/RemoteCustomMetadata.java b/server/src/main/java/org/opensearch/gateway/remote/model/RemoteCustomMetadata.java index ec5dfbec820d4..8e850e903954a 100644 --- a/server/src/main/java/org/opensearch/gateway/remote/model/RemoteCustomMetadata.java +++ b/server/src/main/java/org/opensearch/gateway/remote/model/RemoteCustomMetadata.java @@ -10,7 +10,7 @@ import org.opensearch.cluster.metadata.Metadata.Custom; import org.opensearch.common.io.Streams; -import org.opensearch.common.remote.AbstractRemoteWritableBlobEntity; +import org.opensearch.common.remote.AbstractClusterMetadataWriteableBlobEntity; import org.opensearch.common.remote.BlobPathParameters; import org.opensearch.core.common.io.stream.NamedWriteableAwareStreamInput; import org.opensearch.core.common.io.stream.NamedWriteableRegistry; @@ -32,7 +32,7 @@ /** * Wrapper class for uploading/downloading {@link Custom} to/from remote blob store */ -public class RemoteCustomMetadata extends AbstractRemoteWritableBlobEntity { +public class RemoteCustomMetadata extends AbstractClusterMetadataWriteableBlobEntity { public static final String CUSTOM_METADATA = "custom"; public static final String CUSTOM_DELIMITER = "--"; diff --git a/server/src/main/java/org/opensearch/gateway/remote/model/RemoteDiscoveryNodes.java b/server/src/main/java/org/opensearch/gateway/remote/model/RemoteDiscoveryNodes.java index fb399e2899cdd..446207a767009 100644 --- a/server/src/main/java/org/opensearch/gateway/remote/model/RemoteDiscoveryNodes.java +++ b/server/src/main/java/org/opensearch/gateway/remote/model/RemoteDiscoveryNodes.java @@ -10,7 +10,7 @@ import org.opensearch.cluster.node.DiscoveryNodes; import org.opensearch.common.io.Streams; -import org.opensearch.common.remote.AbstractRemoteWritableBlobEntity; +import org.opensearch.common.remote.AbstractClusterMetadataWriteableBlobEntity; import org.opensearch.common.remote.BlobPathParameters; import org.opensearch.core.compress.Compressor; import org.opensearch.gateway.remote.ClusterMetadataManifest.UploadedMetadata; @@ -29,7 +29,7 @@ /** * Wrapper class for uploading/downloading {@link DiscoveryNodes} to/from remote blob store */ -public class RemoteDiscoveryNodes extends AbstractRemoteWritableBlobEntity { +public class RemoteDiscoveryNodes extends AbstractClusterMetadataWriteableBlobEntity { public static final String DISCOVERY_NODES = "nodes"; public static final ChecksumWritableBlobStoreFormat DISCOVERY_NODES_FORMAT = new ChecksumWritableBlobStoreFormat<>( diff --git a/server/src/main/java/org/opensearch/gateway/remote/model/RemoteGlobalMetadata.java b/server/src/main/java/org/opensearch/gateway/remote/model/RemoteGlobalMetadata.java index 09f07de0d5c24..0082f873f8dba 100644 --- a/server/src/main/java/org/opensearch/gateway/remote/model/RemoteGlobalMetadata.java +++ b/server/src/main/java/org/opensearch/gateway/remote/model/RemoteGlobalMetadata.java @@ -10,7 +10,7 @@ import org.opensearch.cluster.metadata.Metadata; import org.opensearch.common.io.Streams; -import org.opensearch.common.remote.AbstractRemoteWritableBlobEntity; +import org.opensearch.common.remote.AbstractClusterMetadataWriteableBlobEntity; import org.opensearch.common.remote.BlobPathParameters; import org.opensearch.core.compress.Compressor; import org.opensearch.core.xcontent.NamedXContentRegistry; @@ -25,7 +25,7 @@ /** * Wrapper class for uploading/downloading global metadata ({@link Metadata}) to/from remote blob store */ -public class RemoteGlobalMetadata extends AbstractRemoteWritableBlobEntity { +public class RemoteGlobalMetadata extends AbstractClusterMetadataWriteableBlobEntity { public static final String GLOBAL_METADATA = "global_metadata"; public static final ChecksumBlobStoreFormat GLOBAL_METADATA_FORMAT = new ChecksumBlobStoreFormat<>( diff --git a/server/src/main/java/org/opensearch/gateway/remote/model/RemoteHashesOfConsistentSettings.java b/server/src/main/java/org/opensearch/gateway/remote/model/RemoteHashesOfConsistentSettings.java index 1debf75cdfec9..dee48237e5c4c 100644 --- a/server/src/main/java/org/opensearch/gateway/remote/model/RemoteHashesOfConsistentSettings.java +++ b/server/src/main/java/org/opensearch/gateway/remote/model/RemoteHashesOfConsistentSettings.java @@ -10,7 +10,7 @@ import org.opensearch.cluster.metadata.DiffableStringMap; import org.opensearch.common.io.Streams; -import org.opensearch.common.remote.AbstractRemoteWritableBlobEntity; +import org.opensearch.common.remote.AbstractClusterMetadataWriteableBlobEntity; import org.opensearch.common.remote.BlobPathParameters; import org.opensearch.core.compress.Compressor; import org.opensearch.gateway.remote.ClusterMetadataManifest; @@ -28,7 +28,7 @@ /** * Wrapper class for uploading/downloading {@link DiffableStringMap} to/from remote blob store */ -public class RemoteHashesOfConsistentSettings extends AbstractRemoteWritableBlobEntity { +public class RemoteHashesOfConsistentSettings extends AbstractClusterMetadataWriteableBlobEntity { public static final String HASHES_OF_CONSISTENT_SETTINGS = "hashes-of-consistent-settings"; public static final ChecksumWritableBlobStoreFormat HASHES_OF_CONSISTENT_SETTINGS_FORMAT = new ChecksumWritableBlobStoreFormat<>("hashes-of-consistent-settings", DiffableStringMap::readFrom); diff --git a/server/src/main/java/org/opensearch/gateway/remote/model/RemoteIndexMetadata.java b/server/src/main/java/org/opensearch/gateway/remote/model/RemoteIndexMetadata.java index 830b09b92e2cb..5308f92c633b1 100644 --- a/server/src/main/java/org/opensearch/gateway/remote/model/RemoteIndexMetadata.java +++ b/server/src/main/java/org/opensearch/gateway/remote/model/RemoteIndexMetadata.java @@ -10,7 +10,7 @@ import org.opensearch.cluster.metadata.IndexMetadata; import org.opensearch.common.io.Streams; -import org.opensearch.common.remote.AbstractRemoteWritableBlobEntity; +import org.opensearch.common.remote.AbstractClusterMetadataWriteableBlobEntity; import org.opensearch.common.remote.BlobPathParameters; import org.opensearch.core.compress.Compressor; import org.opensearch.core.xcontent.NamedXContentRegistry; @@ -29,7 +29,7 @@ /** * Wrapper class for uploading/downloading {@link IndexMetadata} to/from remote blob store */ -public class RemoteIndexMetadata extends AbstractRemoteWritableBlobEntity { +public class RemoteIndexMetadata extends AbstractClusterMetadataWriteableBlobEntity { public static final int INDEX_METADATA_CURRENT_CODEC_VERSION = 2; diff --git a/server/src/main/java/org/opensearch/gateway/remote/model/RemotePersistentSettingsMetadata.java b/server/src/main/java/org/opensearch/gateway/remote/model/RemotePersistentSettingsMetadata.java index 9ee3db8f289e3..81042f289254c 100644 --- a/server/src/main/java/org/opensearch/gateway/remote/model/RemotePersistentSettingsMetadata.java +++ b/server/src/main/java/org/opensearch/gateway/remote/model/RemotePersistentSettingsMetadata.java @@ -9,7 +9,7 @@ package org.opensearch.gateway.remote.model; import org.opensearch.common.io.Streams; -import org.opensearch.common.remote.AbstractRemoteWritableBlobEntity; +import org.opensearch.common.remote.AbstractClusterMetadataWriteableBlobEntity; import org.opensearch.common.remote.BlobPathParameters; import org.opensearch.common.settings.Settings; import org.opensearch.core.compress.Compressor; @@ -31,7 +31,7 @@ /** * Wrapper class for uploading/downloading persistent {@link Settings} to/from remote blob store */ -public class RemotePersistentSettingsMetadata extends AbstractRemoteWritableBlobEntity { +public class RemotePersistentSettingsMetadata extends AbstractClusterMetadataWriteableBlobEntity { public static final String SETTING_METADATA = "settings"; diff --git a/server/src/main/java/org/opensearch/gateway/remote/model/RemoteRoutingTableBlobStore.java b/server/src/main/java/org/opensearch/gateway/remote/model/RemoteRoutingTableBlobStore.java index 7c4a5bf2236a1..1a28c97dc8a77 100644 --- a/server/src/main/java/org/opensearch/gateway/remote/model/RemoteRoutingTableBlobStore.java +++ b/server/src/main/java/org/opensearch/gateway/remote/model/RemoteRoutingTableBlobStore.java @@ -9,10 +9,13 @@ package org.opensearch.gateway.remote.model; import org.opensearch.common.blobstore.BlobPath; -import org.opensearch.common.remote.AbstractRemoteWritableBlobEntity; +import org.opensearch.common.remote.AbstractClusterMetadataWriteableBlobEntity; +import org.opensearch.common.remote.RemoteWriteableBlobEntity; import org.opensearch.common.remote.RemoteWriteableEntity; +import org.opensearch.common.remote.RemoteWriteableEntityBlobStore; import org.opensearch.common.settings.ClusterSettings; import org.opensearch.common.settings.Setting; +import org.opensearch.gateway.remote.RemoteClusterStateUtils; import org.opensearch.gateway.remote.routingtable.RemoteIndexRoutingTable; import org.opensearch.index.remote.RemoteStoreEnums; import org.opensearch.index.remote.RemoteStorePathStrategy; @@ -28,8 +31,8 @@ * @param which can be uploaded to / downloaded from blob store * @param The concrete class implementing {@link RemoteWriteableEntity} which is used as a wrapper for IndexRoutingTable entity. */ -public class RemoteRoutingTableBlobStore> extends - RemoteClusterStateBlobStore { +public class RemoteRoutingTableBlobStore> extends + RemoteWriteableEntityBlobStore { /** * This setting is used to set the remote routing table store blob store path type strategy. @@ -66,7 +69,14 @@ public RemoteRoutingTableBlobStore( String executor, ClusterSettings clusterSettings ) { - super(blobStoreTransferService, blobStoreRepository, clusterName, threadPool, executor); + super( + blobStoreTransferService, + blobStoreRepository, + clusterName, + threadPool, + executor, + RemoteClusterStateUtils.CLUSTER_STATE_PATH_TOKEN + ); this.pathType = clusterSettings.get(REMOTE_ROUTING_TABLE_PATH_TYPE_SETTING); this.pathHashAlgo = clusterSettings.get(REMOTE_ROUTING_TABLE_PATH_HASH_ALGO_SETTING); clusterSettings.addSettingsUpdateConsumer(REMOTE_ROUTING_TABLE_PATH_TYPE_SETTING, this::setPathTypeSetting); @@ -74,7 +84,7 @@ public RemoteRoutingTableBlobStore( } @Override - public BlobPath getBlobPathForUpload(final AbstractRemoteWritableBlobEntity obj) { + public BlobPath getBlobPathForUpload(final RemoteWriteableBlobEntity obj) { assert obj.getBlobPathParameters().getPathTokens().size() == 1 : "Unexpected tokens in RemoteRoutingTableObject"; BlobPath indexRoutingPath = getBlobPathPrefix(obj.clusterUUID()).add(INDEX_ROUTING_TABLE); diff --git a/server/src/main/java/org/opensearch/gateway/remote/model/RemoteTemplatesMetadata.java b/server/src/main/java/org/opensearch/gateway/remote/model/RemoteTemplatesMetadata.java index 4513d35aef5e8..6ae8a533f9a21 100644 --- a/server/src/main/java/org/opensearch/gateway/remote/model/RemoteTemplatesMetadata.java +++ b/server/src/main/java/org/opensearch/gateway/remote/model/RemoteTemplatesMetadata.java @@ -10,7 +10,7 @@ import org.opensearch.cluster.metadata.TemplatesMetadata; import org.opensearch.common.io.Streams; -import org.opensearch.common.remote.AbstractRemoteWritableBlobEntity; +import org.opensearch.common.remote.AbstractClusterMetadataWriteableBlobEntity; import org.opensearch.common.remote.BlobPathParameters; import org.opensearch.core.compress.Compressor; import org.opensearch.core.xcontent.NamedXContentRegistry; @@ -31,7 +31,7 @@ /** * Wrapper class for uploading/downloading {@link TemplatesMetadata} to/from remote blob store */ -public class RemoteTemplatesMetadata extends AbstractRemoteWritableBlobEntity { +public class RemoteTemplatesMetadata extends AbstractClusterMetadataWriteableBlobEntity { public static final String TEMPLATES_METADATA = "templates"; diff --git a/server/src/main/java/org/opensearch/gateway/remote/model/RemoteTransientSettingsMetadata.java b/server/src/main/java/org/opensearch/gateway/remote/model/RemoteTransientSettingsMetadata.java index fd0526f05d015..d4f3837f80084 100644 --- a/server/src/main/java/org/opensearch/gateway/remote/model/RemoteTransientSettingsMetadata.java +++ b/server/src/main/java/org/opensearch/gateway/remote/model/RemoteTransientSettingsMetadata.java @@ -9,7 +9,7 @@ package org.opensearch.gateway.remote.model; import org.opensearch.common.io.Streams; -import org.opensearch.common.remote.AbstractRemoteWritableBlobEntity; +import org.opensearch.common.remote.AbstractClusterMetadataWriteableBlobEntity; import org.opensearch.common.remote.BlobPathParameters; import org.opensearch.common.settings.Settings; import org.opensearch.core.compress.Compressor; @@ -32,7 +32,7 @@ /** * Wrapper class for uploading/downloading transient {@link Settings} to/from remote blob store */ -public class RemoteTransientSettingsMetadata extends AbstractRemoteWritableBlobEntity { +public class RemoteTransientSettingsMetadata extends AbstractClusterMetadataWriteableBlobEntity { public static final String TRANSIENT_SETTING_METADATA = "transient-settings"; diff --git a/server/src/main/java/org/opensearch/gateway/remote/routingtable/RemoteIndexRoutingTable.java b/server/src/main/java/org/opensearch/gateway/remote/routingtable/RemoteIndexRoutingTable.java index 40b5bafde2b13..46c5074c48eb8 100644 --- a/server/src/main/java/org/opensearch/gateway/remote/routingtable/RemoteIndexRoutingTable.java +++ b/server/src/main/java/org/opensearch/gateway/remote/routingtable/RemoteIndexRoutingTable.java @@ -10,7 +10,7 @@ import org.opensearch.cluster.routing.IndexRoutingTable; import org.opensearch.common.io.Streams; -import org.opensearch.common.remote.AbstractRemoteWritableBlobEntity; +import org.opensearch.common.remote.AbstractClusterMetadataWriteableBlobEntity; import org.opensearch.common.remote.BlobPathParameters; import org.opensearch.core.compress.Compressor; import org.opensearch.core.index.Index; @@ -27,7 +27,7 @@ /** * Remote store object for IndexRoutingTable */ -public class RemoteIndexRoutingTable extends AbstractRemoteWritableBlobEntity { +public class RemoteIndexRoutingTable extends AbstractClusterMetadataWriteableBlobEntity { public static final String INDEX_ROUTING_TABLE = "index-routing"; public static final String INDEX_ROUTING_METADATA_PREFIX = "indexRouting--"; diff --git a/server/src/main/java/org/opensearch/gateway/remote/routingtable/RemoteRoutingTableDiff.java b/server/src/main/java/org/opensearch/gateway/remote/routingtable/RemoteRoutingTableDiff.java index e876d939490d0..2370417dc14df 100644 --- a/server/src/main/java/org/opensearch/gateway/remote/routingtable/RemoteRoutingTableDiff.java +++ b/server/src/main/java/org/opensearch/gateway/remote/routingtable/RemoteRoutingTableDiff.java @@ -12,7 +12,7 @@ import org.opensearch.cluster.routing.IndexRoutingTable; import org.opensearch.cluster.routing.RoutingTableIncrementalDiff; import org.opensearch.common.io.Streams; -import org.opensearch.common.remote.AbstractRemoteWritableBlobEntity; +import org.opensearch.common.remote.AbstractClusterMetadataWriteableBlobEntity; import org.opensearch.common.remote.BlobPathParameters; import org.opensearch.core.compress.Compressor; import org.opensearch.gateway.remote.ClusterMetadataManifest; @@ -30,7 +30,7 @@ * Represents a incremental difference between {@link org.opensearch.cluster.routing.RoutingTable} objects that can be serialized and deserialized. * This class is responsible for writing and reading the differences between RoutingTables to and from an input/output stream. */ -public class RemoteRoutingTableDiff extends AbstractRemoteWritableBlobEntity { +public class RemoteRoutingTableDiff extends AbstractClusterMetadataWriteableBlobEntity { private final RoutingTableIncrementalDiff routingTableIncrementalDiff; private long term; diff --git a/server/src/test/java/org/opensearch/common/remote/AbstractRemoteWritableEntityManagerTests.java b/server/src/test/java/org/opensearch/common/remote/AbstractRemoteWritableEntityManagerTests.java index 3d10bbf59f6ad..d236274c99236 100644 --- a/server/src/test/java/org/opensearch/common/remote/AbstractRemoteWritableEntityManagerTests.java +++ b/server/src/test/java/org/opensearch/common/remote/AbstractRemoteWritableEntityManagerTests.java @@ -24,7 +24,7 @@ public void testGetStoreWithKnownEntityType() { String knownEntityType = "knownType"; RemoteWritableEntityStore mockStore = mock(RemoteWritableEntityStore.class); manager.remoteWritableEntityStores.put(knownEntityType, mockStore); - AbstractRemoteWritableBlobEntity mockEntity = mock(AbstractRemoteWritableBlobEntity.class); + AbstractClusterMetadataWriteableBlobEntity mockEntity = mock(AbstractClusterMetadataWriteableBlobEntity.class); when(mockEntity.getType()).thenReturn(knownEntityType); RemoteWritableEntityStore store = manager.getStore(mockEntity); @@ -35,7 +35,7 @@ public void testGetStoreWithKnownEntityType() { public void testGetStoreWithUnknownEntityType() { AbstractRemoteWritableEntityManager manager = new ConcreteRemoteWritableEntityManager(); String unknownEntityType = "unknownType"; - AbstractRemoteWritableBlobEntity mockEntity = mock(AbstractRemoteWritableBlobEntity.class); + AbstractClusterMetadataWriteableBlobEntity mockEntity = mock(AbstractClusterMetadataWriteableBlobEntity.class); when(mockEntity.getType()).thenReturn(unknownEntityType); assertThrows(IllegalArgumentException.class, () -> manager.getStore(mockEntity)); @@ -46,7 +46,7 @@ private static class ConcreteRemoteWritableEntityManager extends AbstractRemoteW @Override protected ActionListener getWrappedWriteListener( String component, - AbstractRemoteWritableBlobEntity remoteEntity, + AbstractClusterMetadataWriteableBlobEntity remoteEntity, ActionListener listener ) { return null; @@ -55,7 +55,7 @@ protected ActionListener getWrappedWriteListener( @Override protected ActionListener getWrappedReadListener( String component, - AbstractRemoteWritableBlobEntity remoteEntity, + AbstractClusterMetadataWriteableBlobEntity remoteEntity, ActionListener listener ) { return null; diff --git a/server/src/test/java/org/opensearch/gateway/remote/RemoteClusterStateServiceTests.java b/server/src/test/java/org/opensearch/gateway/remote/RemoteClusterStateServiceTests.java index 59ca62dff2aa7..1871f4d08ba43 100644 --- a/server/src/test/java/org/opensearch/gateway/remote/RemoteClusterStateServiceTests.java +++ b/server/src/test/java/org/opensearch/gateway/remote/RemoteClusterStateServiceTests.java @@ -42,7 +42,7 @@ import org.opensearch.common.compress.DeflateCompressor; import org.opensearch.common.lucene.store.ByteArrayIndexInput; import org.opensearch.common.network.NetworkModule; -import org.opensearch.common.remote.AbstractRemoteWritableBlobEntity; +import org.opensearch.common.remote.AbstractClusterMetadataWriteableBlobEntity; import org.opensearch.common.settings.ClusterSettings; import org.opensearch.common.settings.Settings; import org.opensearch.common.util.FeatureFlags; @@ -3464,7 +3464,7 @@ public static DiscoveryNodes nodesWithLocalNodeClusterManager() { return DiscoveryNodes.builder().clusterManagerNodeId("cluster-manager-id").localNodeId("cluster-manager-id").add(localNode).build(); } - private class BlobNameMatcher implements ArgumentMatcher { + private class BlobNameMatcher implements ArgumentMatcher { private final String expectedBlobName; BlobNameMatcher(String expectedBlobName) { @@ -3472,7 +3472,7 @@ private class BlobNameMatcher implements ArgumentMatcher Date: Tue, 13 Aug 2024 08:33:52 -0700 Subject: [PATCH 24/40] Improve breaking change logic for changelog verifier (#15214) Signed-off-by: Jay Deng --- .github/workflows/changelog_verifier.yml | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/.github/workflows/changelog_verifier.yml b/.github/workflows/changelog_verifier.yml index cf9343c2c3aac..cd0415119282c 100644 --- a/.github/workflows/changelog_verifier.yml +++ b/.github/workflows/changelog_verifier.yml @@ -27,18 +27,29 @@ jobs: continue-on-error: true - run: | # The check was possibly skipped leading to success for both the jobs + has_backport_label=${{ contains(join(github.event.pull_request.labels.*.name, ', '), 'backport')}} + has_breaking_label=${{ contains(join(github.event.pull_request.labels.*.name, ', '), '>breaking')}} + if [[ $has_breaking_label == true && $has_backport_label == true ]]; then + echo "error: Please make sure that the PR does not have a backport label associated with it when making breaking changes" + exit 1 + fi + if [[ ${{ steps.verify-changelog-3x.outcome }} == 'success' && ${{ steps.verify-changelog.outcome }} == 'success' ]]; then exit 0 fi - + if [[ ${{ steps.verify-changelog-3x.outcome }} == 'failure' && ${{ steps.verify-changelog.outcome }} == 'failure' ]]; then echo "error: Please ensure a changelog entry exists in CHANGELOG.md or CHANGELOG-3.0.md" exit 1 fi - + # Concatenates the labels and checks if the string contains "backport" - has_backport_label=${{ contains(join(github.event.pull_request.labels.*.name, ', '), 'backport')}} if [[ ${{ steps.verify-changelog.outcome }} == 'success' && $has_backport_label == false ]]; then echo "error: Please make sure that the PR has a backport label associated with it when making an entry to the CHANGELOG.md file" exit 1 fi + + if [[ ${{ steps.verify-changelog-3x.outcome }} == 'success' && $has_backport_label == true ]]; then + echo "error: Please make sure that the PR does not have a backport label associated with it when making an entry to the CHANGELOG-3.0.md file" + exit 1 + fi From ea1ab7ca8c1183577e7a250b9d7006dc0d709cac Mon Sep 17 00:00:00 2001 From: 10000-ki <10000ki6472@gmail.com> Date: Wed, 14 Aug 2024 01:11:41 +0900 Subject: [PATCH 25/40] Correction of a typo (#15220) Signed-off-by: 10000-ki <10000ki6472@gmail.com> --- .../org/opensearch/core/xcontent/AbstractXContentParser.java | 2 +- .../test/java/org/opensearch/packaging/util/FileUtils.java | 2 +- .../test/java/org/opensearch/http/DanglingIndicesRestIT.java | 4 ++-- .../org/opensearch/indices/recovery/DanglingIndicesIT.java | 4 ++-- server/src/main/java/org/opensearch/http/HttpChannel.java | 2 +- .../opensearch/indices/recovery/RecoverySourceHandler.java | 2 +- .../repositories/blobstore/BlobStoreRepository.java | 2 +- .../src/main/java/org/opensearch/snapshots/package-info.java | 2 +- server/src/main/java/org/opensearch/transport/TcpChannel.java | 2 +- .../action/admin/cluster/node/info/NodeInfoTests.java | 2 +- .../test/java/org/opensearch/common/LocalTimeOffsetTests.java | 2 +- .../indices/breaker/HierarchyCircuitBreakerServiceTests.java | 2 +- 12 files changed, 14 insertions(+), 14 deletions(-) diff --git a/libs/core/src/main/java/org/opensearch/core/xcontent/AbstractXContentParser.java b/libs/core/src/main/java/org/opensearch/core/xcontent/AbstractXContentParser.java index 4efaacecd0e67..4605aa684db1c 100644 --- a/libs/core/src/main/java/org/opensearch/core/xcontent/AbstractXContentParser.java +++ b/libs/core/src/main/java/org/opensearch/core/xcontent/AbstractXContentParser.java @@ -375,7 +375,7 @@ private static void skipToListStart(XContentParser parser) throws IOException { } } - // read a list without bounds checks, assuming the the current parser is always on an array start + // read a list without bounds checks, assuming the current parser is always on an array start private static List readListUnsafe(XContentParser parser, Supplier> mapFactory) throws IOException { assert parser.currentToken() == Token.START_ARRAY; ArrayList list = new ArrayList<>(); diff --git a/qa/os/src/test/java/org/opensearch/packaging/util/FileUtils.java b/qa/os/src/test/java/org/opensearch/packaging/util/FileUtils.java index 5169ce18fff79..dd5248738569e 100644 --- a/qa/os/src/test/java/org/opensearch/packaging/util/FileUtils.java +++ b/qa/os/src/test/java/org/opensearch/packaging/util/FileUtils.java @@ -380,7 +380,7 @@ public static String escapePath(Path path) { } /** - * Recursively copy the the source directory to the target directory, preserving permissions. + * Recursively copy the source directory to the target directory, preserving permissions. */ public static void copyDirectory(Path source, Path target) throws IOException { Files.walkFileTree(source, new SimpleFileVisitor() { diff --git a/qa/smoke-test-http/src/test/java/org/opensearch/http/DanglingIndicesRestIT.java b/qa/smoke-test-http/src/test/java/org/opensearch/http/DanglingIndicesRestIT.java index 42c7fd667fd8f..741660f972bfb 100644 --- a/qa/smoke-test-http/src/test/java/org/opensearch/http/DanglingIndicesRestIT.java +++ b/qa/smoke-test-http/src/test/java/org/opensearch/http/DanglingIndicesRestIT.java @@ -152,8 +152,8 @@ public void testDanglingIndicesCanBeImported() throws Exception { * 1, then create two indices and delete them both while one node in * the cluster is stopped. The deletion of the second pushes the deletion * of the first out of the graveyard. When the stopped node is resumed, - * only the second index will be found into the graveyard and the the - * other will be considered dangling, and can therefore be listed and + * only the second index will be found into the graveyard and the other + * will be considered dangling, and can therefore be listed and * deleted through the API */ public void testDanglingIndicesCanBeDeleted() throws Exception { diff --git a/server/src/internalClusterTest/java/org/opensearch/indices/recovery/DanglingIndicesIT.java b/server/src/internalClusterTest/java/org/opensearch/indices/recovery/DanglingIndicesIT.java index 8fd7961cab3a7..7bdf33edf1534 100644 --- a/server/src/internalClusterTest/java/org/opensearch/indices/recovery/DanglingIndicesIT.java +++ b/server/src/internalClusterTest/java/org/opensearch/indices/recovery/DanglingIndicesIT.java @@ -298,8 +298,8 @@ public void testMustAcceptDataLossToImportDanglingIndex() throws Exception { * 1, then create two indices and delete them both while one node in * the cluster is stopped. The deletion of the second pushes the deletion * of the first out of the graveyard. When the stopped node is resumed, - * only the second index will be found into the graveyard and the the - * other will be considered dangling, and can therefore be listed and + * only the second index will be found into the graveyard and the other + * will be considered dangling, and can therefore be listed and * deleted through the API */ public void testDanglingIndexCanBeDeleted() throws Exception { diff --git a/server/src/main/java/org/opensearch/http/HttpChannel.java b/server/src/main/java/org/opensearch/http/HttpChannel.java index ed20ec89a9099..7048f08faff9f 100644 --- a/server/src/main/java/org/opensearch/http/HttpChannel.java +++ b/server/src/main/java/org/opensearch/http/HttpChannel.java @@ -77,7 +77,7 @@ default void handleException(Exception ex) {} /** * Returns the contextual property associated with this specific HTTP channel (the - * implementation of how such properties are managed depends on the the particular + * implementation of how such properties are managed depends on the particular * transport engine). * * @param name the name of the property diff --git a/server/src/main/java/org/opensearch/indices/recovery/RecoverySourceHandler.java b/server/src/main/java/org/opensearch/indices/recovery/RecoverySourceHandler.java index abf9b1aaeb2cc..8966516d3ef7f 100644 --- a/server/src/main/java/org/opensearch/indices/recovery/RecoverySourceHandler.java +++ b/server/src/main/java/org/opensearch/indices/recovery/RecoverySourceHandler.java @@ -540,7 +540,7 @@ void sendFiles(Store store, StoreFileMetadata[] files, IntSupplier translogOps, void createRetentionLease(final long startingSeqNo, ActionListener listener) { RunUnderPrimaryPermit.run(() -> { - // Clone the peer recovery retention lease belonging to the source shard. We are retaining history between the the local + // Clone the peer recovery retention lease belonging to the source shard. We are retaining history between the local // checkpoint of the safe commit we're creating and this lease's retained seqno with the retention lock, and by cloning an // existing lease we (approximately) know that all our peers are also retaining history as requested by the cloned lease. If // the recovery now fails before copying enough history over then a subsequent attempt will find this lease, determine it is diff --git a/server/src/main/java/org/opensearch/repositories/blobstore/BlobStoreRepository.java b/server/src/main/java/org/opensearch/repositories/blobstore/BlobStoreRepository.java index c4908f8c5fc4b..3e6a75565891f 100644 --- a/server/src/main/java/org/opensearch/repositories/blobstore/BlobStoreRepository.java +++ b/server/src/main/java/org/opensearch/repositories/blobstore/BlobStoreRepository.java @@ -400,7 +400,7 @@ public abstract class BlobStoreRepository extends AbstractLifecycleComponent imp /** * Flag that is set to {@code true} if this instance is started with {@link #metadata} that has a higher value for * {@link RepositoryMetadata#pendingGeneration()} than for {@link RepositoryMetadata#generation()} indicating a full cluster restart - * potentially accounting for the the last {@code index-N} write in the cluster state. + * potentially accounting for the last {@code index-N} write in the cluster state. * Note: While it is true that this value could also be set to {@code true} for an instance on a node that is just joining the cluster * during a new {@code index-N} write, this does not present a problem. The node will still load the correct {@link RepositoryData} in * all cases and simply do a redundant listing of the repository contents if it tries to load {@link RepositoryData} and falls back diff --git a/server/src/main/java/org/opensearch/snapshots/package-info.java b/server/src/main/java/org/opensearch/snapshots/package-info.java index a573d7a136620..02fbf737dbf62 100644 --- a/server/src/main/java/org/opensearch/snapshots/package-info.java +++ b/server/src/main/java/org/opensearch/snapshots/package-info.java @@ -131,7 +131,7 @@ * snapshots, we load the {@link org.opensearch.snapshots.SnapshotInfo} for the source snapshot and check for shard snapshot * failures of the relevant indices. *
  • Once all shard counts are known and the health of all source indices data has been verified, we populate the - * {@code SnapshotsInProgress.Entry#clones} map for the clone operation with the the relevant shard clone tasks.
    • + * {@code SnapshotsInProgress.Entry#clones} map for the clone operation with the relevant shard clone tasks. *
  • After the clone tasks have been added to the {@code SnapshotsInProgress.Entry}, cluster-manager executes them on its snapshot thread-pool * by invoking {@link org.opensearch.repositories.Repository#cloneShardSnapshot} for each shard that is to be cloned. Each completed * shard snapshot triggers a call to the {@link org.opensearch.snapshots.SnapshotsService#SHARD_STATE_EXECUTOR} which updates the diff --git a/server/src/main/java/org/opensearch/transport/TcpChannel.java b/server/src/main/java/org/opensearch/transport/TcpChannel.java index 7d4515de85d80..75a6d8b2cff5f 100644 --- a/server/src/main/java/org/opensearch/transport/TcpChannel.java +++ b/server/src/main/java/org/opensearch/transport/TcpChannel.java @@ -101,7 +101,7 @@ public interface TcpChannel extends CloseableChannel { /** * Returns the contextual property associated with this specific TCP channel (the - * implementation of how such properties are managed depends on the the particular + * implementation of how such properties are managed depends on the particular * transport engine). * * @param name the name of the property diff --git a/server/src/test/java/org/opensearch/action/admin/cluster/node/info/NodeInfoTests.java b/server/src/test/java/org/opensearch/action/admin/cluster/node/info/NodeInfoTests.java index cfd6fcec4bdc6..b372e9c4dac5b 100644 --- a/server/src/test/java/org/opensearch/action/admin/cluster/node/info/NodeInfoTests.java +++ b/server/src/test/java/org/opensearch/action/admin/cluster/node/info/NodeInfoTests.java @@ -52,7 +52,7 @@ public class NodeInfoTests extends OpenSearchTestCase { /** - * Check that the the {@link NodeInfo#getInfo(Class)} method returns null + * Check that the {@link NodeInfo#getInfo(Class)} method returns null * for absent info objects, and returns the right thing for present info * objects. */ diff --git a/server/src/test/java/org/opensearch/common/LocalTimeOffsetTests.java b/server/src/test/java/org/opensearch/common/LocalTimeOffsetTests.java index b032e27397f2d..cfef82b2d65a1 100644 --- a/server/src/test/java/org/opensearch/common/LocalTimeOffsetTests.java +++ b/server/src/test/java/org/opensearch/common/LocalTimeOffsetTests.java @@ -302,7 +302,7 @@ private static long time(String time, ZoneId zone) { } /** - * The the last "fully defined" transitions in the provided {@linkplain ZoneId}. + * The last "fully defined" transitions in the provided {@linkplain ZoneId}. */ private static ZoneOffsetTransition lastTransitionIn(ZoneId zone) { List transitions = zone.getRules().getTransitions(); diff --git a/server/src/test/java/org/opensearch/indices/breaker/HierarchyCircuitBreakerServiceTests.java b/server/src/test/java/org/opensearch/indices/breaker/HierarchyCircuitBreakerServiceTests.java index 591f8f15a1441..ea113abd6913a 100644 --- a/server/src/test/java/org/opensearch/indices/breaker/HierarchyCircuitBreakerServiceTests.java +++ b/server/src/test/java/org/opensearch/indices/breaker/HierarchyCircuitBreakerServiceTests.java @@ -793,7 +793,7 @@ public void testAllocationBucketsBreaker() { // make sure used bytes is greater than the total circuit breaker limit breaker.addWithoutBreaking(200); - // make sure that we check on the the following call + // make sure that we check on the following call for (int i = 0; i < 1023; i++) { multiBucketConsumer.accept(0); } From 54c13a6ae6afcdb18a8ffb1b1ef8044ec3b1ce56 Mon Sep 17 00:00:00 2001 From: Derek Ho Date: Tue, 13 Aug 2024 12:48:36 -0400 Subject: [PATCH 26/40] Bump dnsjava version to 3.6.0 to address CVE-2024-25638 (#15173) * Bump dnsjava version Signed-off-by: Derek Ho * Change to implementation Signed-off-by: Derek Ho * Update test/fixtures/hdfs-fixture/build.gradle Co-authored-by: Andriy Redko Signed-off-by: Derek Ho --------- Signed-off-by: Derek Ho Signed-off-by: Derek Ho Co-authored-by: Andriy Redko --- test/fixtures/hdfs-fixture/build.gradle | 1 + 1 file changed, 1 insertion(+) diff --git a/test/fixtures/hdfs-fixture/build.gradle b/test/fixtures/hdfs-fixture/build.gradle index aaa150e73a13e..adf420cf84663 100644 --- a/test/fixtures/hdfs-fixture/build.gradle +++ b/test/fixtures/hdfs-fixture/build.gradle @@ -55,6 +55,7 @@ dependencies { exclude group: 'com.nimbusds' exclude module: "commons-configuration2" } + api "dnsjava:dnsjava:3.6.0" api "org.codehaus.jettison:jettison:${versions.jettison}" api "org.apache.commons:commons-compress:${versions.commonscompress}" api "commons-codec:commons-codec:${versions.commonscodec}" From 625dd5a3bcd1a56d07f49492225413eae5db58fb Mon Sep 17 00:00:00 2001 From: Rishab Nahata Date: Wed, 14 Aug 2024 01:20:11 +0530 Subject: [PATCH 27/40] Fix responsibility check for existing shards allocator when timed out (#15223) * Fix responsibility check for existing shards allocator when timed out Signed-off-by: Rishab Nahata --- .../gateway/BaseGatewayShardAllocator.java | 8 +++++ .../gateway/PrimaryShardAllocator.java | 2 +- .../gateway/ReplicaShardAllocator.java | 2 +- .../gateway/ReplicaShardBatchAllocator.java | 2 +- .../PrimaryShardBatchAllocatorTests.java | 33 +++---------------- .../ReplicaShardBatchAllocatorTests.java | 18 ++++++++++ 6 files changed, 34 insertions(+), 31 deletions(-) diff --git a/server/src/main/java/org/opensearch/gateway/BaseGatewayShardAllocator.java b/server/src/main/java/org/opensearch/gateway/BaseGatewayShardAllocator.java index 41704545c7a6f..2b6c5e3f5ae53 100644 --- a/server/src/main/java/org/opensearch/gateway/BaseGatewayShardAllocator.java +++ b/server/src/main/java/org/opensearch/gateway/BaseGatewayShardAllocator.java @@ -90,12 +90,20 @@ protected void allocateUnassignedBatchOnTimeout(Set shardIds, RoutingAl ShardRouting unassignedShard = iterator.next(); AllocateUnassignedDecision allocationDecision; if (unassignedShard.primary() == primary && shardIds.contains(unassignedShard.shardId())) { + if (isResponsibleFor(unassignedShard) == false) { + continue; + } allocationDecision = AllocateUnassignedDecision.throttle(null); executeDecision(unassignedShard, allocationDecision, allocation, iterator); } } } + /** + * Is the allocator responsible for allocating the given {@link ShardRouting}? + */ + protected abstract boolean isResponsibleFor(ShardRouting shardRouting); + protected void executeDecision( ShardRouting shardRouting, AllocateUnassignedDecision allocateUnassignedDecision, diff --git a/server/src/main/java/org/opensearch/gateway/PrimaryShardAllocator.java b/server/src/main/java/org/opensearch/gateway/PrimaryShardAllocator.java index f41545cbdf9bf..dea7ca9a08edd 100644 --- a/server/src/main/java/org/opensearch/gateway/PrimaryShardAllocator.java +++ b/server/src/main/java/org/opensearch/gateway/PrimaryShardAllocator.java @@ -82,7 +82,7 @@ public abstract class PrimaryShardAllocator extends BaseGatewayShardAllocator { /** * Is the allocator responsible for allocating the given {@link ShardRouting}? */ - protected static boolean isResponsibleFor(final ShardRouting shard) { + protected boolean isResponsibleFor(final ShardRouting shard) { return shard.primary() // must be primary && shard.unassigned() // must be unassigned // only handle either an existing store or a snapshot recovery diff --git a/server/src/main/java/org/opensearch/gateway/ReplicaShardAllocator.java b/server/src/main/java/org/opensearch/gateway/ReplicaShardAllocator.java index aaf0d696e1444..c30ee8479ac97 100644 --- a/server/src/main/java/org/opensearch/gateway/ReplicaShardAllocator.java +++ b/server/src/main/java/org/opensearch/gateway/ReplicaShardAllocator.java @@ -191,7 +191,7 @@ public void processExistingRecoveries(RoutingAllocation allocation) { /** * Is the allocator responsible for allocating the given {@link ShardRouting}? */ - protected static boolean isResponsibleFor(final ShardRouting shard) { + protected boolean isResponsibleFor(final ShardRouting shard) { return shard.primary() == false // must be a replica && shard.unassigned() // must be unassigned // if we are allocating a replica because of index creation, no need to go and find a copy, there isn't one... diff --git a/server/src/main/java/org/opensearch/gateway/ReplicaShardBatchAllocator.java b/server/src/main/java/org/opensearch/gateway/ReplicaShardBatchAllocator.java index 0818b187271cb..020a543ac5fc5 100644 --- a/server/src/main/java/org/opensearch/gateway/ReplicaShardBatchAllocator.java +++ b/server/src/main/java/org/opensearch/gateway/ReplicaShardBatchAllocator.java @@ -173,7 +173,7 @@ private AllocateUnassignedDecision getUnassignedShardAllocationDecision( RoutingAllocation allocation, Supplier> nodeStoreFileMetaDataMapSupplier ) { - if (!isResponsibleFor(shardRouting)) { + if (isResponsibleFor(shardRouting) == false) { return AllocateUnassignedDecision.NOT_TAKEN; } Tuple> result = canBeAllocatedToAtLeastOneNode(shardRouting, allocation); diff --git a/server/src/test/java/org/opensearch/gateway/PrimaryShardBatchAllocatorTests.java b/server/src/test/java/org/opensearch/gateway/PrimaryShardBatchAllocatorTests.java index 48183fed66671..2edde8281b11a 100644 --- a/server/src/test/java/org/opensearch/gateway/PrimaryShardBatchAllocatorTests.java +++ b/server/src/test/java/org/opensearch/gateway/PrimaryShardBatchAllocatorTests.java @@ -51,6 +51,7 @@ import java.util.stream.Collectors; import static org.opensearch.cluster.routing.UnassignedInfo.Reason.CLUSTER_RECOVERED; +import static org.opensearch.cluster.routing.UnassignedInfo.Reason.INDEX_CREATED; public class PrimaryShardBatchAllocatorTests extends OpenSearchAllocationTestCase { @@ -283,40 +284,16 @@ public void testAllocateUnassignedBatchOnTimeoutWithNoMatchingPrimaryShards() { assertEquals(0, ignoredShards.size()); } - public void testAllocateUnassignedBatchOnTimeoutWithNonPrimaryShards() { + public void testAllocateUnassignedBatchOnTimeoutSkipIgnoringNewPrimaryShards() { ClusterSettings clusterSettings = new ClusterSettings(Settings.EMPTY, ClusterSettings.BUILT_IN_CLUSTER_SETTINGS); AllocationDeciders allocationDeciders = randomAllocationDeciders(Settings.builder().build(), clusterSettings, random()); setUpShards(1); - final RoutingAllocation routingAllocation = routingAllocationWithOnePrimary(allocationDeciders, CLUSTER_RECOVERED, "allocId-0"); + final RoutingAllocation routingAllocation = routingAllocationWithOnePrimary(allocationDeciders, INDEX_CREATED); + ShardRouting shardRouting = routingAllocation.routingTable().getIndicesRouting().get("test").shard(shardId.id()).primaryShard(); - ShardRouting shardRouting = routingAllocation.routingTable() - .getIndicesRouting() - .get("test") - .shard(shardId.id()) - .replicaShards() - .get(0); Set shardIds = new HashSet<>(); shardIds.add(shardRouting.shardId()); - batchAllocator.allocateUnassignedBatchOnTimeout(shardIds, routingAllocation, false); - - List ignoredShards = routingAllocation.routingNodes().unassigned().ignored(); - assertEquals(1, ignoredShards.size()); - } - - public void testAllocateUnassignedBatchOnTimeoutWithNoShards() { - ClusterSettings clusterSettings = new ClusterSettings(Settings.EMPTY, ClusterSettings.BUILT_IN_CLUSTER_SETTINGS); - AllocationDeciders allocationDeciders = randomAllocationDeciders(Settings.builder().build(), clusterSettings, random()); - setUpShards(1); - final RoutingAllocation routingAllocation = routingAllocationWithOnePrimary(allocationDeciders, CLUSTER_RECOVERED, "allocId-0"); - - ShardRouting shardRouting = routingAllocation.routingTable() - .getIndicesRouting() - .get("test") - .shard(shardId.id()) - .replicaShards() - .get(0); - Set shardIds = new HashSet<>(); - batchAllocator.allocateUnassignedBatchOnTimeout(shardIds, routingAllocation, false); + batchAllocator.allocateUnassignedBatchOnTimeout(shardIds, routingAllocation, true); List ignoredShards = routingAllocation.routingNodes().unassigned().ignored(); assertEquals(0, ignoredShards.size()); diff --git a/server/src/test/java/org/opensearch/gateway/ReplicaShardBatchAllocatorTests.java b/server/src/test/java/org/opensearch/gateway/ReplicaShardBatchAllocatorTests.java index 78ed3f2c7d38c..988723e023a2a 100644 --- a/server/src/test/java/org/opensearch/gateway/ReplicaShardBatchAllocatorTests.java +++ b/server/src/test/java/org/opensearch/gateway/ReplicaShardBatchAllocatorTests.java @@ -744,6 +744,24 @@ public void testAllocateUnassignedBatchOnTimeoutWithAlreadyRecoveringReplicaShar assertThat(allocation.routingNodes().unassigned().ignored().size(), equalTo(0)); } + public void testAllocateUnassignedBatchOnTimeoutSkipIgnoringNewReplicaShards() { + RoutingAllocation allocation = onePrimaryOnNode1And1Replica( + yesAllocationDeciders(), + Settings.EMPTY, + UnassignedInfo.Reason.INDEX_CREATED + ); + final RoutingNodes.UnassignedShards.UnassignedIterator iterator = allocation.routingNodes().unassigned().iterator(); + Set shards = new HashSet<>(); + while (iterator.hasNext()) { + ShardRouting sr = iterator.next(); + if (sr.primary() == false) { + shards.add(sr.shardId()); + } + } + testBatchAllocator.allocateUnassignedBatchOnTimeout(shards, allocation, false); + assertThat(allocation.routingNodes().unassigned().ignored().size(), equalTo(0)); + } + private RoutingAllocation onePrimaryOnNode1And1Replica(AllocationDeciders deciders) { return onePrimaryOnNode1And1Replica(deciders, Settings.EMPTY, UnassignedInfo.Reason.CLUSTER_RECOVERED); } From 6e27e5ad8e37944605bef61e14f1972f440e4724 Mon Sep 17 00:00:00 2001 From: Naomichi Yamakita Date: Wed, 14 Aug 2024 06:25:57 +0900 Subject: [PATCH 28/40] Fix parseToken array popping (#13620) * Fix parseToken array popping Signed-off-by: naomichi-y * Add fix description Signed-off-by: naomichi-y * Refactor JsonToStringXContentParser#parseToken This method was pretty complicated, hard to follow, and therefore prone to bugs. This change introduces a proper stack (rather than a StringBuilder) to keep track of fields under fields. When ever we parse a field name, we push it onto the stack, recursively parse its value (an object, an array, or a primitive value), then we pop (guaranteeing pushes and pops are balanced). Signed-off-by: Michael Froh * Remove use of org.apache.logging.log4j.util.Strings Signed-off-by: Michael Froh * Make sure JsonToStringXContentParser ends on END_OBJECT There is logic in DocumentParser that expects any object parser to start with currentToken() pointing to START_OBJECT, and return with currentToken() pointing to END_OBJECT. My previous commits were stepping over the start/end, and returning on the token following the end. Signed-off-by: Michael Froh * Throw exception when flat object is null The existing behavior throws an exception on null flat object by advancing the parser beyond the end of the flat object input. We could simply skip a null flat_object field, but this would be a behavioral change from 2.7. Signed-off-by: Michael Froh --------- Signed-off-by: naomichi-y Signed-off-by: Michael Froh Co-authored-by: Michael Froh --- CHANGELOG.md | 1 + .../xcontent/JsonToStringXContentParser.java | 122 ++++++++---------- .../index/mapper/FlatObjectFieldMapper.java | 9 +- .../JsonToStringXContentParserTests.java | 57 +++++++- 4 files changed, 117 insertions(+), 72 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 34cd4c2097e48..cf0d032cc75ea 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -46,6 +46,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), - Fix missing value of FieldSort for unsigned_long ([#14963](https://github.com/opensearch-project/OpenSearch/pull/14963)) - Fix delete index template failed when the index template matches a data stream but is unused ([#15080](https://github.com/opensearch-project/OpenSearch/pull/15080)) - Fix array_index_out_of_bounds_exception when indexing documents with field name containing only dot ([#15126](https://github.com/opensearch-project/OpenSearch/pull/15126)) +- Fixed array field name omission in flat_object function for nested JSON ([#13620](https://github.com/opensearch-project/OpenSearch/pull/13620)) ### Security diff --git a/server/src/main/java/org/opensearch/common/xcontent/JsonToStringXContentParser.java b/server/src/main/java/org/opensearch/common/xcontent/JsonToStringXContentParser.java index 998122d9e5c43..d24571fc5778d 100644 --- a/server/src/main/java/org/opensearch/common/xcontent/JsonToStringXContentParser.java +++ b/server/src/main/java/org/opensearch/common/xcontent/JsonToStringXContentParser.java @@ -9,6 +9,7 @@ package org.opensearch.common.xcontent; import org.opensearch.common.xcontent.json.JsonXContent; +import org.opensearch.core.common.Strings; import org.opensearch.core.common.bytes.BytesReference; import org.opensearch.core.xcontent.AbstractXContentParser; import org.opensearch.core.xcontent.DeprecationHandler; @@ -23,6 +24,9 @@ import java.math.BigInteger; import java.nio.CharBuffer; import java.util.ArrayList; +import java.util.Collections; +import java.util.Deque; +import java.util.LinkedList; /** * JsonToStringParser is the main parser class to transform JSON into stringFields in a XContentParser @@ -32,21 +36,20 @@ */ public class JsonToStringXContentParser extends AbstractXContentParser { private final String fieldTypeName; - private XContentParser parser; + private final XContentParser parser; - private ArrayList valueList = new ArrayList<>(); - private ArrayList valueAndPathList = new ArrayList<>(); - private ArrayList keyList = new ArrayList<>(); + private final ArrayList valueList = new ArrayList<>(); + private final ArrayList valueAndPathList = new ArrayList<>(); + private final ArrayList keyList = new ArrayList<>(); - private XContentBuilder builder = XContentBuilder.builder(JsonXContent.jsonXContent); + private final XContentBuilder builder = XContentBuilder.builder(JsonXContent.jsonXContent); - private NamedXContentRegistry xContentRegistry; + private final NamedXContentRegistry xContentRegistry; - private DeprecationHandler deprecationHandler; + private final DeprecationHandler deprecationHandler; private static final String VALUE_AND_PATH_SUFFIX = "._valueAndPath"; private static final String VALUE_SUFFIX = "._value"; - private static final String DOT_SYMBOL = "."; private static final String EQUAL_SYMBOL = "="; public JsonToStringXContentParser( @@ -63,9 +66,14 @@ public JsonToStringXContentParser( } public XContentParser parseObject() throws IOException { + assert currentToken() == Token.START_OBJECT; + parser.nextToken(); // Skip the outer START_OBJECT. Need to return on END_OBJECT. + builder.startObject(); - StringBuilder path = new StringBuilder(fieldTypeName); - parseToken(path, null); + LinkedList path = new LinkedList<>(Collections.singleton(fieldTypeName)); + while (currentToken() != Token.END_OBJECT) { + parseToken(path); + } builder.field(this.fieldTypeName, keyList); builder.field(this.fieldTypeName + VALUE_SUFFIX, valueList); builder.field(this.fieldTypeName + VALUE_AND_PATH_SUFFIX, valueAndPathList); @@ -74,75 +82,55 @@ public XContentParser parseObject() throws IOException { return JsonXContent.jsonXContent.createParser(this.xContentRegistry, this.deprecationHandler, String.valueOf(jString)); } - private void parseToken(StringBuilder path, String currentFieldName) throws IOException { - - while (this.parser.nextToken() != Token.END_OBJECT) { - if (this.parser.currentName() != null) { - currentFieldName = this.parser.currentName(); + private void parseToken(Deque path) throws IOException { + if (this.parser.currentToken() == Token.FIELD_NAME) { + String fieldName = this.parser.currentName(); + path.addLast(fieldName); // Pushing onto the stack *must* be matched by pop + String parts = fieldName; + while (parts.contains(".")) { // Extract the intermediate keys maybe present in fieldName + int dotPos = parts.indexOf('.'); + String part = parts.substring(0, dotPos); + this.keyList.add(part); + parts = parts.substring(dotPos + 1); } - StringBuilder parsedFields = new StringBuilder(); - - if (this.parser.currentToken() == Token.FIELD_NAME) { - path.append(DOT_SYMBOL).append(currentFieldName); - int dotIndex = currentFieldName.indexOf(DOT_SYMBOL); - String fieldNameSuffix = currentFieldName; - // The field name may be of the form foo.bar.baz - // If that's the case, each "part" is a key. - while (dotIndex >= 0) { - String fieldNamePrefix = fieldNameSuffix.substring(0, dotIndex); - if (!fieldNamePrefix.isEmpty()) { - this.keyList.add(fieldNamePrefix); - } - fieldNameSuffix = fieldNameSuffix.substring(dotIndex + 1); - dotIndex = fieldNameSuffix.indexOf(DOT_SYMBOL); - } - if (!fieldNameSuffix.isEmpty()) { - this.keyList.add(fieldNameSuffix); - } - } else if (this.parser.currentToken() == Token.START_ARRAY) { - parseToken(path, currentFieldName); - break; - } else if (this.parser.currentToken() == Token.END_ARRAY) { - // skip - } else if (this.parser.currentToken() == Token.START_OBJECT) { - parseToken(path, currentFieldName); - int dotIndex = path.lastIndexOf(DOT_SYMBOL, path.length()); - - if (dotIndex != -1 && path.length() > currentFieldName.length()) { - path.setLength(path.length() - currentFieldName.length() - 1); - } - } else { - if (!path.toString().contains(currentFieldName)) { - path.append(DOT_SYMBOL).append(currentFieldName); - } - parseValue(parsedFields); - this.valueList.add(parsedFields.toString()); - this.valueAndPathList.add(path + EQUAL_SYMBOL + parsedFields); - int dotIndex = path.lastIndexOf(DOT_SYMBOL, path.length()); - if (dotIndex != -1 && path.length() > currentFieldName.length()) { - path.setLength(path.length() - currentFieldName.length() - 1); - } + this.keyList.add(parts); // parts has no dot, so either it's the original fieldName or it's the last part + this.parser.nextToken(); // advance to the value of fieldName + parseToken(path); // parse the value for fieldName (which will be an array, an object, or a primitive value) + path.removeLast(); // Here is where we pop fieldName from the stack (since we're done with the value of fieldName) + // Note that whichever other branch we just passed through has already ended with nextToken(), so we + // don't need to call it. + } else if (this.parser.currentToken() == Token.START_ARRAY) { + parser.nextToken(); + while (this.parser.currentToken() != Token.END_ARRAY) { + parseToken(path); } - + this.parser.nextToken(); + } else if (this.parser.currentToken() == Token.START_OBJECT) { + parser.nextToken(); + while (this.parser.currentToken() != Token.END_OBJECT) { + parseToken(path); + } + this.parser.nextToken(); + } else if (this.parser.currentToken().isValue()) { + String parsedValue = parseValue(); + if (parsedValue != null) { + this.valueList.add(parsedValue); + this.valueAndPathList.add(Strings.collectionToDelimitedString(path, ".") + EQUAL_SYMBOL + parsedValue); + } + this.parser.nextToken(); } } - private void parseValue(StringBuilder parsedFields) throws IOException { + private String parseValue() throws IOException { switch (this.parser.currentToken()) { case VALUE_BOOLEAN: case VALUE_NUMBER: case VALUE_STRING: case VALUE_NULL: - parsedFields.append(this.parser.textOrNull()); - break; + return this.parser.textOrNull(); // Handle other token types as needed - case FIELD_NAME: - case VALUE_EMBEDDED_OBJECT: - case END_ARRAY: - case START_ARRAY: - break; default: - throw new IOException("Unsupported token type [" + parser.currentToken() + "]"); + throw new IOException("Unsupported value token type [" + parser.currentToken() + "]"); } } diff --git a/server/src/main/java/org/opensearch/index/mapper/FlatObjectFieldMapper.java b/server/src/main/java/org/opensearch/index/mapper/FlatObjectFieldMapper.java index 9a3f2595a7c9e..b82fa3999612a 100644 --- a/server/src/main/java/org/opensearch/index/mapper/FlatObjectFieldMapper.java +++ b/server/src/main/java/org/opensearch/index/mapper/FlatObjectFieldMapper.java @@ -568,8 +568,12 @@ protected void parseCreateField(ParseContext context) throws IOException { if (context.externalValueSet()) { String value = context.externalValue().toString(); parseValueAddFields(context, value, fieldType().name()); + } else if (context.parser().currentToken() == XContentParser.Token.VALUE_NULL) { + context.parser().nextToken(); // This triggers an exception in DocumentParser. + // We could remove the above nextToken() call to skip the null value, but the existing + // behavior (since 2.7) throws the exception. } else { - JsonToStringXContentParser JsonToStringParser = new JsonToStringXContentParser( + JsonToStringXContentParser jsonToStringParser = new JsonToStringXContentParser( NamedXContentRegistry.EMPTY, DeprecationHandler.IGNORE_DEPRECATIONS, context.parser(), @@ -579,7 +583,7 @@ protected void parseCreateField(ParseContext context) throws IOException { JsonToStringParser is the main parser class to transform JSON into stringFields in a XContentParser It reads the JSON object and parsed to a list of string */ - XContentParser parser = JsonToStringParser.parseObject(); + XContentParser parser = jsonToStringParser.parseObject(); XContentParser.Token currentToken; while ((currentToken = parser.nextToken()) != XContentParser.Token.END_OBJECT) { @@ -594,7 +598,6 @@ protected void parseCreateField(ParseContext context) throws IOException { } } - } } diff --git a/server/src/test/java/org/opensearch/common/xcontent/JsonToStringXContentParserTests.java b/server/src/test/java/org/opensearch/common/xcontent/JsonToStringXContentParserTests.java index 0feb7bcd1ceec..a0f5150981a08 100644 --- a/server/src/test/java/org/opensearch/common/xcontent/JsonToStringXContentParserTests.java +++ b/server/src/test/java/org/opensearch/common/xcontent/JsonToStringXContentParserTests.java @@ -19,7 +19,6 @@ public class JsonToStringXContentParserTests extends OpenSearchTestCase { private String flattenJsonString(String fieldName, String in) throws IOException { - String transformed; try ( XContentParser parser = JsonXContent.jsonXContent.createParser( xContentRegistry(), @@ -33,10 +32,11 @@ private String flattenJsonString(String fieldName, String in) throws IOException parser, fieldName ); - // Skip the START_OBJECT token: + // Point to the first token (should be START_OBJECT) jsonToStringXContentParser.nextToken(); XContentParser transformedParser = jsonToStringXContentParser.parseObject(); + assertSame(XContentParser.Token.END_OBJECT, jsonToStringXContentParser.currentToken()); try (XContentBuilder jsonBuilder = XContentFactory.jsonBuilder()) { jsonBuilder.copyCurrentStructure(transformedParser); return jsonBuilder.toString(); @@ -110,4 +110,57 @@ public void testNestChildObjectWithDotsAndFieldWithDots() throws IOException { ); } + public void testArrayOfObjects() throws IOException { + String jsonExample = "{" + + "\"field\": {" + + " \"detail\": {" + + " \"foooooooooooo\": [" + + " {\"name\":\"baz\"}," + + " {\"name\":\"baz\"}" + + " ]" + + " }" + + "}}"; + + assertEquals( + "{" + + "\"flat\":[\"field\",\"detail\",\"foooooooooooo\",\"name\",\"name\"]," + + "\"flat._value\":[\"baz\",\"baz\"]," + + "\"flat._valueAndPath\":[" + + "\"flat.field.detail.foooooooooooo.name=baz\"," + + "\"flat.field.detail.foooooooooooo.name=baz\"" + + "]}", + flattenJsonString("flat", jsonExample) + ); + } + + public void testArraysOfObjectsAndValues() throws IOException { + String jsonExample = "{" + + "\"field\": {" + + " \"detail\": {" + + " \"foooooooooooo\": [" + + " {\"name\":\"baz\"}," + + " {\"name\":\"baz\"}" + + " ]" + + " }," + + " \"numbers\" : [" + + " 1," + + " 2," + + " 3" + + " ]" + + "}}"; + + assertEquals( + "{" + + "\"flat\":[\"field\",\"detail\",\"foooooooooooo\",\"name\",\"name\",\"numbers\"]," + + "\"flat._value\":[\"baz\",\"baz\",\"1\",\"2\",\"3\"]," + + "\"flat._valueAndPath\":[" + + "\"flat.field.detail.foooooooooooo.name=baz\"," + + "\"flat.field.detail.foooooooooooo.name=baz\"," + + "\"flat.field.numbers=1\"," + + "\"flat.field.numbers=2\"," + + "\"flat.field.numbers=3\"" + + "]}", + flattenJsonString("flat", jsonExample) + ); + } } From f1f721f55ef727a69f4f0f180096461de5a94406 Mon Sep 17 00:00:00 2001 From: Rishabh Singh Date: Tue, 13 Aug 2024 19:20:57 -0700 Subject: [PATCH 29/40] Fix logic for getting maintainer ids for benchmark workflow (#15236) --- .github/workflows/benchmark-pull-request.yml | 20 ++++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/.github/workflows/benchmark-pull-request.yml b/.github/workflows/benchmark-pull-request.yml index 59b3e3900823e..c494df6e27ce3 100644 --- a/.github/workflows/benchmark-pull-request.yml +++ b/.github/workflows/benchmark-pull-request.yml @@ -108,13 +108,25 @@ jobs: echo "prHeadRepo=$headRepo" >> $GITHUB_ENV echo "prHeadRefSha=$headRefSha" >> $GITHUB_ENV - id: get_approvers - run: | - echo "approvers=$(cat .github/CODEOWNERS | grep '^\*' | tr -d '* ' | sed 's/@/,/g' | sed 's/,//1')" >> $GITHUB_OUTPUT + uses: actions/github-script@v7 + with: + github-token: ${{ secrets.GITHUB_TOKEN }} + result-encoding: string + script: | + // Get the collaborators - filtered to maintainer permissions + const maintainersResponse = await github.request('GET /repos/{owner}/{repo}/collaborators', { + owner: context.repo.owner, + repo: context.repo.repo, + permission: 'maintain', + affiliation: 'all', + per_page: 100 + }); + return maintainersResponse.data.map(item => item.login).join(', '); - uses: trstringer/manual-approval@v1 - if: (!contains(steps.get_approvers.outputs.approvers, github.event.comment.user.login)) + if: (!contains(steps.get_approvers.outputs.result, github.event.comment.user.login)) with: secret: ${{ github.TOKEN }} - approvers: ${{ steps.get_approvers.outputs.approvers }} + approvers: ${{ steps.get_approvers.outputs.result }} minimum-approvals: 1 issue-title: 'Request to approve/deny benchmark run for PR #${{ env.PR_NUMBER }}' issue-body: "Please approve or deny the benchmark run for PR #${{ env.PR_NUMBER }}" From 68bdb774c8e71690e863e84998d28918f667ac04 Mon Sep 17 00:00:00 2001 From: Sooraj Sinha <81695996+soosinha@users.noreply.github.com> Date: Wed, 14 Aug 2024 14:24:19 +0530 Subject: [PATCH 30/40] Use clusterUUID for discovery nodes remote path (#15143) Signed-off-by: Sooraj Sinha --- .../remote/RemoteStatePublicationIT.java | 155 ++++++++++++++++++ .../remote/RemoteClusterStateService.java | 2 +- 2 files changed, 156 insertions(+), 1 deletion(-) create mode 100644 server/src/internalClusterTest/java/org/opensearch/gateway/remote/RemoteStatePublicationIT.java diff --git a/server/src/internalClusterTest/java/org/opensearch/gateway/remote/RemoteStatePublicationIT.java b/server/src/internalClusterTest/java/org/opensearch/gateway/remote/RemoteStatePublicationIT.java new file mode 100644 index 0000000000000..07d6e1379ced8 --- /dev/null +++ b/server/src/internalClusterTest/java/org/opensearch/gateway/remote/RemoteStatePublicationIT.java @@ -0,0 +1,155 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.gateway.remote; + +import org.opensearch.action.admin.cluster.settings.ClusterUpdateSettingsRequest; +import org.opensearch.action.admin.cluster.state.ClusterStateResponse; +import org.opensearch.client.Client; +import org.opensearch.common.blobstore.BlobPath; +import org.opensearch.common.settings.Settings; +import org.opensearch.common.util.FeatureFlags; +import org.opensearch.gateway.remote.model.RemoteClusterMetadataManifest; +import org.opensearch.indices.recovery.RecoverySettings; +import org.opensearch.remotestore.RemoteStoreBaseIntegTestCase; +import org.opensearch.repositories.RepositoriesService; +import org.opensearch.repositories.blobstore.BlobStoreRepository; +import org.opensearch.repositories.fs.ReloadableFsRepository; +import org.opensearch.test.OpenSearchIntegTestCase.ClusterScope; +import org.opensearch.test.OpenSearchIntegTestCase.Scope; +import org.junit.Before; + +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.util.Base64; +import java.util.Locale; +import java.util.Map; +import java.util.function.Function; +import java.util.stream.Collectors; + +import static org.opensearch.gateway.remote.RemoteClusterStateAttributesManager.DISCOVERY_NODES; +import static org.opensearch.gateway.remote.RemoteClusterStateService.REMOTE_CLUSTER_STATE_ENABLED_SETTING; +import static org.opensearch.gateway.remote.RemoteClusterStateUtils.DELIMITER; +import static org.opensearch.gateway.remote.model.RemoteClusterBlocks.CLUSTER_BLOCKS; +import static org.opensearch.gateway.remote.model.RemoteCoordinationMetadata.COORDINATION_METADATA; +import static org.opensearch.gateway.remote.model.RemoteCustomMetadata.CUSTOM_METADATA; +import static org.opensearch.gateway.remote.model.RemotePersistentSettingsMetadata.SETTING_METADATA; +import static org.opensearch.gateway.remote.model.RemoteTemplatesMetadata.TEMPLATES_METADATA; +import static org.opensearch.gateway.remote.model.RemoteTransientSettingsMetadata.TRANSIENT_SETTING_METADATA; +import static org.opensearch.node.remotestore.RemoteStoreNodeAttribute.REMOTE_STORE_REPOSITORY_SETTINGS_ATTRIBUTE_KEY_PREFIX; +import static org.opensearch.node.remotestore.RemoteStoreNodeAttribute.REMOTE_STORE_REPOSITORY_TYPE_ATTRIBUTE_KEY_FORMAT; +import static org.opensearch.node.remotestore.RemoteStoreNodeAttribute.REMOTE_STORE_ROUTING_TABLE_REPOSITORY_NAME_ATTRIBUTE_KEY; +import static org.opensearch.test.hamcrest.OpenSearchAssertions.assertAcked; + +@ClusterScope(scope = Scope.TEST, numDataNodes = 0) +public class RemoteStatePublicationIT extends RemoteStoreBaseIntegTestCase { + + private static String INDEX_NAME = "test-index"; + + @Before + public void setup() { + asyncUploadMockFsRepo = false; + } + + @Override + protected Settings featureFlagSettings() { + return Settings.builder().put(super.featureFlagSettings()).put(FeatureFlags.REMOTE_PUBLICATION_EXPERIMENTAL, "true").build(); + } + + @Override + protected Settings nodeSettings(int nodeOrdinal) { + String routingTableRepoName = "remote-routing-repo"; + String routingTableRepoTypeAttributeKey = String.format( + Locale.getDefault(), + "node.attr." + REMOTE_STORE_REPOSITORY_TYPE_ATTRIBUTE_KEY_FORMAT, + routingTableRepoName + ); + String routingTableRepoSettingsAttributeKeyPrefix = String.format( + Locale.getDefault(), + "node.attr." + REMOTE_STORE_REPOSITORY_SETTINGS_ATTRIBUTE_KEY_PREFIX, + routingTableRepoName + ); + return Settings.builder() + .put(super.nodeSettings(nodeOrdinal)) + .put(REMOTE_CLUSTER_STATE_ENABLED_SETTING.getKey(), true) + .put("node.attr." + REMOTE_STORE_ROUTING_TABLE_REPOSITORY_NAME_ATTRIBUTE_KEY, routingTableRepoName) + .put(routingTableRepoTypeAttributeKey, ReloadableFsRepository.TYPE) + .put(routingTableRepoSettingsAttributeKeyPrefix + "location", segmentRepoPath) + .build(); + } + + public void testPublication() throws Exception { + // create cluster with multi node (3 master + 2 data) + prepareCluster(3, 2, INDEX_NAME, 1, 2); + ensureStableCluster(5); + ensureGreen(INDEX_NAME); + // update settings on a random node + assertAcked( + internalCluster().client() + .admin() + .cluster() + .updateSettings( + new ClusterUpdateSettingsRequest().persistentSettings( + Settings.builder().put(RecoverySettings.INDICES_RECOVERY_MAX_BYTES_PER_SEC_SETTING.getKey(), "10mb").build() + ) + ) + .actionGet() + ); + + RepositoriesService repositoriesService = internalCluster().getClusterManagerNodeInstance(RepositoriesService.class); + BlobStoreRepository repository = (BlobStoreRepository) repositoriesService.repository(REPOSITORY_NAME); + + Map globalMetadataFiles = getMetadataFiles(repository, RemoteClusterStateUtils.GLOBAL_METADATA_PATH_TOKEN); + + assertTrue(globalMetadataFiles.containsKey(COORDINATION_METADATA)); + assertTrue(globalMetadataFiles.containsKey(SETTING_METADATA)); + assertTrue(globalMetadataFiles.containsKey(TRANSIENT_SETTING_METADATA)); + assertTrue(globalMetadataFiles.containsKey(TEMPLATES_METADATA)); + assertTrue(globalMetadataFiles.keySet().stream().anyMatch(key -> key.startsWith(CUSTOM_METADATA))); + + Map ephemeralMetadataFiles = getMetadataFiles( + repository, + RemoteClusterStateUtils.CLUSTER_STATE_EPHEMERAL_PATH_TOKEN + ); + + assertTrue(ephemeralMetadataFiles.containsKey(CLUSTER_BLOCKS)); + assertTrue(ephemeralMetadataFiles.containsKey(DISCOVERY_NODES)); + + Map manifestFiles = getMetadataFiles(repository, RemoteClusterMetadataManifest.MANIFEST); + assertTrue(manifestFiles.containsKey(RemoteClusterMetadataManifest.MANIFEST)); + + // get settings from each node and verify that it is updated + Settings settings = clusterService().getSettings(); + logger.info("settings : {}", settings); + for (Client client : clients()) { + ClusterStateResponse response = client.admin().cluster().prepareState().clear().setMetadata(true).get(); + String refreshSetting = response.getState() + .metadata() + .settings() + .get(RecoverySettings.INDICES_RECOVERY_MAX_BYTES_PER_SEC_SETTING.getKey()); + assertEquals("10mb", refreshSetting); + } + } + + private Map getMetadataFiles(BlobStoreRepository repository, String subDirectory) throws IOException { + BlobPath metadataPath = repository.basePath() + .add( + Base64.getUrlEncoder() + .withoutPadding() + .encodeToString(getClusterState().getClusterName().value().getBytes(StandardCharsets.UTF_8)) + ) + .add(RemoteClusterStateUtils.CLUSTER_STATE_PATH_TOKEN) + .add(getClusterState().metadata().clusterUUID()) + .add(subDirectory); + return repository.blobStore().blobContainer(metadataPath).listBlobs().keySet().stream().map(fileName -> { + logger.info(fileName); + return fileName.split(DELIMITER)[0]; + }).collect(Collectors.toMap(Function.identity(), key -> 1, Integer::sum)); + } + +} diff --git a/server/src/main/java/org/opensearch/gateway/remote/RemoteClusterStateService.java b/server/src/main/java/org/opensearch/gateway/remote/RemoteClusterStateService.java index 674279f2251bd..910f601a81ca8 100644 --- a/server/src/main/java/org/opensearch/gateway/remote/RemoteClusterStateService.java +++ b/server/src/main/java/org/opensearch/gateway/remote/RemoteClusterStateService.java @@ -603,7 +603,7 @@ UploadedMetadataResults writeMetadataInParallel( new RemoteDiscoveryNodes( clusterState.nodes(), clusterState.version(), - clusterState.stateUUID(), + clusterState.metadata().clusterUUID(), blobStoreRepository.getCompressor() ), listener From ef1a79fdc7159852a6fc13e61625fa6e661fc6d1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Luk=C3=A1=C5=A1=20Vl=C4=8Dek?= Date: Wed, 14 Aug 2024 11:48:20 +0200 Subject: [PATCH 31/40] Don't rely on test code execution time span for RemoteSegmentTransferTrackerTests (#15187) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Lukáš Vlček --- .../remote/RemoteSegmentTransferTracker.java | 4 ++-- .../remote/RemoteSegmentTransferTrackerTests.java | 15 +++++++++------ 2 files changed, 11 insertions(+), 8 deletions(-) diff --git a/server/src/main/java/org/opensearch/index/remote/RemoteSegmentTransferTracker.java b/server/src/main/java/org/opensearch/index/remote/RemoteSegmentTransferTracker.java index f1843ea3eef38..a29bd1d840b43 100644 --- a/server/src/main/java/org/opensearch/index/remote/RemoteSegmentTransferTracker.java +++ b/server/src/main/java/org/opensearch/index/remote/RemoteSegmentTransferTracker.java @@ -65,7 +65,7 @@ public class RemoteSegmentTransferTracker extends RemoteTransferTracker { private volatile long remoteRefreshSeqNo; /** - * The refresh time of most recent remote refresh. + * The refresh time of the most recent remote refresh. */ private volatile long remoteRefreshTimeMs; @@ -76,7 +76,7 @@ public class RemoteSegmentTransferTracker extends RemoteTransferTracker { private volatile long remoteRefreshStartTimeMs = -1; /** - * The refresh time(clock) of most recent remote refresh. + * The refresh time(clock) of the most recent remote refresh. */ private volatile long remoteRefreshClockTimeMs; diff --git a/server/src/test/java/org/opensearch/index/remote/RemoteSegmentTransferTrackerTests.java b/server/src/test/java/org/opensearch/index/remote/RemoteSegmentTransferTrackerTests.java index 1ec1e9977a9d5..f4101cb054687 100644 --- a/server/src/test/java/org/opensearch/index/remote/RemoteSegmentTransferTrackerTests.java +++ b/server/src/test/java/org/opensearch/index/remote/RemoteSegmentTransferTrackerTests.java @@ -152,15 +152,18 @@ public void testComputeTimeLagOnUpdate() throws InterruptedException { transferTracker.updateLocalRefreshTimeMs(currentTimeMsUsingSystemNanos()); transferTracker.updateLatestLocalFileNameLengthMap(List.of("test"), k -> 1L); - // Sleep for 100ms and then the lag should be within 100ms +/- 20ms - Thread.sleep(100); - assertTrue(Math.abs(transferTracker.getTimeMsLag() - 100) <= 20); + // Sleep for 100ms and then the lag should not be shorter + long span = 100; + Thread.sleep(span); + long lag = transferTracker.getTimeMsLag(); + assertTrue("Actual lag [" + lag + "ms] is not expected to be shorter than span [" + span + "ms]", lag >= span); transferTracker.updateRemoteRefreshTimeMs(transferTracker.getLocalRefreshTimeMs()); transferTracker.updateLocalRefreshTimeMs(currentTimeMsUsingSystemNanos()); - long random = randomIntBetween(50, 200); - Thread.sleep(random); - assertTrue(Math.abs(transferTracker.getTimeMsLag() - random) <= 20); + long randomSpan = randomIntBetween(50, 200); + Thread.sleep(randomSpan); + lag = transferTracker.getTimeMsLag(); + assertTrue("Actual lag [" + lag + "ms] is not expected to be shorter than span [" + randomSpan + "ms]", lag >= randomSpan); } public void testAddUploadBytesStarted() { From d74e276af68edbc5caf382fd2e61f11b9935c221 Mon Sep 17 00:00:00 2001 From: Craig Perkins Date: Wed, 14 Aug 2024 08:15:22 -0400 Subject: [PATCH 32/40] Replace and block usages of org.apache.logging.log4j.util.Strings (#15238) * Replace and block usages of org.apache.logging.log4j.util.Strings Signed-off-by: Craig Perkins * Add to CHANGELOG Signed-off-by: Craig Perkins --------- Signed-off-by: Craig Perkins --- CHANGELOG.md | 1 + .../main/resources/forbidden/opensearch-all-signatures.txt | 3 +++ .../common/logging/JsonThrowablePatternConverter.java | 4 ++-- .../src/main/java/org/opensearch/env/NodeEnvironment.java | 4 ++-- .../gateway/remote/RemoteClusterStateCleanupManager.java | 6 +++--- server/src/main/java/org/opensearch/index/IndexModule.java | 5 ++--- .../src/main/java/org/opensearch/index/shard/ShardPath.java | 4 ++-- .../org/opensearch/client/RestClientBuilderTestCase.java | 3 +-- 8 files changed, 16 insertions(+), 14 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index cf0d032cc75ea..cd0af3932371d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -35,6 +35,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), ### Changed - Add lower limit for primary and replica batch allocators timeout ([#14979](https://github.com/opensearch-project/OpenSearch/pull/14979)) +- Replace and block usages of org.apache.logging.log4j.util.Strings ([#15238](https://github.com/opensearch-project/OpenSearch/pull/15238)) ### Deprecated diff --git a/buildSrc/src/main/resources/forbidden/opensearch-all-signatures.txt b/buildSrc/src/main/resources/forbidden/opensearch-all-signatures.txt index f9f24fd1e2367..199e206450178 100644 --- a/buildSrc/src/main/resources/forbidden/opensearch-all-signatures.txt +++ b/buildSrc/src/main/resources/forbidden/opensearch-all-signatures.txt @@ -17,6 +17,9 @@ java.nio.file.Paths @ Use org.opensearch.common.io.PathUtils.get() instead. java.nio.file.FileSystems#getDefault() @ use org.opensearch.common.io.PathUtils.getDefaultFileSystem() instead. +joptsimple.internal.Strings @ use org.opensearch.core.common.Strings instead. +org.apache.logging.log4j.util.Strings @ use org.opensearch.core.common.Strings instead. + java.nio.file.Files#getFileStore(java.nio.file.Path) @ Use org.opensearch.env.Environment.getFileStore() instead, impacted by JDK-8034057 java.nio.file.Files#isWritable(java.nio.file.Path) @ Use org.opensearch.env.Environment.isWritable() instead, impacted by JDK-8034057 diff --git a/server/src/main/java/org/opensearch/common/logging/JsonThrowablePatternConverter.java b/server/src/main/java/org/opensearch/common/logging/JsonThrowablePatternConverter.java index ed324e4e62d8f..ee21c343e2ea1 100644 --- a/server/src/main/java/org/opensearch/common/logging/JsonThrowablePatternConverter.java +++ b/server/src/main/java/org/opensearch/common/logging/JsonThrowablePatternConverter.java @@ -38,7 +38,7 @@ import org.apache.logging.log4j.core.pattern.ExtendedThrowablePatternConverter; import org.apache.logging.log4j.core.pattern.PatternConverter; import org.apache.logging.log4j.core.pattern.ThrowablePatternConverter; -import org.apache.logging.log4j.util.Strings; +import org.opensearch.core.common.Strings; import java.nio.charset.Charset; import java.util.StringJoiner; @@ -84,7 +84,7 @@ public static JsonThrowablePatternConverter newInstance(final Configuration conf @Override public void format(final LogEvent event, final StringBuilder toAppendTo) { String consoleStacktrace = formatStacktrace(event); - if (Strings.isNotEmpty(consoleStacktrace)) { + if (!Strings.isNullOrEmpty(consoleStacktrace)) { String jsonStacktrace = formatJson(consoleStacktrace); toAppendTo.append(", "); diff --git a/server/src/main/java/org/opensearch/env/NodeEnvironment.java b/server/src/main/java/org/opensearch/env/NodeEnvironment.java index 2748938d8b761..709c0eba4f57f 100644 --- a/server/src/main/java/org/opensearch/env/NodeEnvironment.java +++ b/server/src/main/java/org/opensearch/env/NodeEnvironment.java @@ -35,7 +35,6 @@ import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; import org.apache.logging.log4j.message.ParameterizedMessage; -import org.apache.logging.log4j.util.Strings; import org.apache.lucene.index.IndexWriter; import org.apache.lucene.index.SegmentInfos; import org.apache.lucene.store.Directory; @@ -61,6 +60,7 @@ import org.opensearch.common.settings.Settings; import org.opensearch.common.unit.TimeValue; import org.opensearch.common.util.io.IOUtils; +import org.opensearch.core.common.Strings; import org.opensearch.core.common.unit.ByteSizeValue; import org.opensearch.core.index.Index; import org.opensearch.core.index.shard.ShardId; @@ -1300,7 +1300,7 @@ public static List collectFileCacheDataPath(NodePath fileCacheNodePath) th * Resolve the custom path for a index's shard. */ public static Path resolveBaseCustomLocation(String customDataPath, Path sharedDataPath, int nodeLockId) { - if (Strings.isNotEmpty(customDataPath)) { + if (!Strings.isNullOrEmpty(customDataPath)) { // This assert is because this should be caught by MetadataCreateIndexService assert sharedDataPath != null; return sharedDataPath.resolve(customDataPath).resolve(Integer.toString(nodeLockId)); diff --git a/server/src/main/java/org/opensearch/gateway/remote/RemoteClusterStateCleanupManager.java b/server/src/main/java/org/opensearch/gateway/remote/RemoteClusterStateCleanupManager.java index 8691187c7fbfa..02db15477ff95 100644 --- a/server/src/main/java/org/opensearch/gateway/remote/RemoteClusterStateCleanupManager.java +++ b/server/src/main/java/org/opensearch/gateway/remote/RemoteClusterStateCleanupManager.java @@ -11,7 +11,6 @@ import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; import org.apache.logging.log4j.message.ParameterizedMessage; -import org.apache.logging.log4j.util.Strings; import org.opensearch.cluster.ClusterState; import org.opensearch.cluster.routing.remote.RemoteRoutingTableService; import org.opensearch.cluster.service.ClusterApplierService; @@ -23,6 +22,7 @@ import org.opensearch.common.unit.TimeValue; import org.opensearch.common.util.concurrent.AbstractAsyncTask; import org.opensearch.core.action.ActionListener; +import org.opensearch.core.common.Strings; import org.opensearch.index.translog.transfer.BlobStoreTransferService; import org.opensearch.threadpool.ThreadPool; @@ -125,8 +125,8 @@ void cleanUpStaleFiles() { ClusterState currentAppliedState = clusterApplierService.state(); if (currentAppliedState.nodes().isLocalNodeElectedClusterManager()) { long cleanUpAttemptStateVersion = currentAppliedState.version(); - assert Strings.isNotEmpty(currentAppliedState.getClusterName().value()) : "cluster name is not set"; - assert Strings.isNotEmpty(currentAppliedState.metadata().clusterUUID()) : "cluster uuid is not set"; + assert !Strings.isNullOrEmpty(currentAppliedState.getClusterName().value()) : "cluster name is not set"; + assert !Strings.isNullOrEmpty(currentAppliedState.metadata().clusterUUID()) : "cluster uuid is not set"; if (cleanUpAttemptStateVersion - lastCleanupAttemptStateVersion > SKIP_CLEANUP_STATE_CHANGES) { logger.info( "Cleaning up stale remote state files for cluster [{}] with uuid [{}]. Last clean was done before {} updates", diff --git a/server/src/main/java/org/opensearch/index/IndexModule.java b/server/src/main/java/org/opensearch/index/IndexModule.java index eab070e1c6c10..dc1bf94662385 100644 --- a/server/src/main/java/org/opensearch/index/IndexModule.java +++ b/server/src/main/java/org/opensearch/index/IndexModule.java @@ -98,6 +98,7 @@ import java.util.HashMap; import java.util.HashSet; import java.util.List; +import java.util.Locale; import java.util.Map; import java.util.Objects; import java.util.Set; @@ -108,8 +109,6 @@ import java.util.function.Function; import java.util.function.Supplier; -import static org.apache.logging.log4j.util.Strings.toRootUpperCase; - /** * IndexModule represents the central extension point for index level custom implementations like: *
      @@ -580,7 +579,7 @@ public enum DataLocalityType { public static DataLocalityType getValueOf(final String localityType) { Objects.requireNonNull(localityType, "No locality type given."); - final String localityTypeName = toRootUpperCase(localityType.trim()); + final String localityTypeName = localityType.trim().toUpperCase(Locale.ROOT); final DataLocalityType type = LOCALITY_TYPES.get(localityTypeName); if (type != null) { return type; diff --git a/server/src/main/java/org/opensearch/index/shard/ShardPath.java b/server/src/main/java/org/opensearch/index/shard/ShardPath.java index d16e5707b2e0e..911bfec94e190 100644 --- a/server/src/main/java/org/opensearch/index/shard/ShardPath.java +++ b/server/src/main/java/org/opensearch/index/shard/ShardPath.java @@ -32,10 +32,10 @@ package org.opensearch.index.shard; import org.apache.logging.log4j.Logger; -import org.apache.logging.log4j.util.Strings; import org.opensearch.cluster.metadata.IndexMetadata; import org.opensearch.common.annotation.PublicApi; import org.opensearch.common.util.io.IOUtils; +import org.opensearch.core.common.Strings; import org.opensearch.core.index.shard.ShardId; import org.opensearch.core.xcontent.NamedXContentRegistry; import org.opensearch.env.NodeEnvironment; @@ -205,7 +205,7 @@ public static ShardPath loadShardPath( } else { final Path dataPath; final Path statePath = loadedPath; - final boolean hasCustomDataPath = Strings.isNotEmpty(customDataPath); + final boolean hasCustomDataPath = !Strings.isNullOrEmpty(customDataPath); if (hasCustomDataPath) { dataPath = NodeEnvironment.resolveCustomLocation(customDataPath, shardId, sharedDataPath, nodeLockId); } else { diff --git a/test/framework/src/main/java/org/opensearch/client/RestClientBuilderTestCase.java b/test/framework/src/main/java/org/opensearch/client/RestClientBuilderTestCase.java index 5a964954c4599..757d04041492b 100644 --- a/test/framework/src/main/java/org/opensearch/client/RestClientBuilderTestCase.java +++ b/test/framework/src/main/java/org/opensearch/client/RestClientBuilderTestCase.java @@ -32,7 +32,6 @@ package org.opensearch.client; -import joptsimple.internal.Strings; import org.apache.hc.core5.http.Header; import org.opensearch.test.OpenSearchTestCase; @@ -55,7 +54,7 @@ public void assertHeaders(RestClient client, Map expectedHeaders assertEquals(expectedValue, header.getValue()); } if (expectedHeaders.isEmpty() == false) { - fail("Missing expected headers in rest client: " + Strings.join(expectedHeaders.keySet(), ", ")); + fail("Missing expected headers in rest client: " + String.join(",", expectedHeaders.keySet())); } } } From 68b74ff20d92167c928365cba07aee8e333e9a05 Mon Sep 17 00:00:00 2001 From: Varun Bansal Date: Wed, 14 Aug 2024 18:57:09 +0530 Subject: [PATCH 33/40] [Offline Nodes] Adds new library for offline tasks (#13574) --------- Signed-off-by: Varun Bansal Signed-off-by: Bukhtawar Khan Co-authored-by: Bukhtawar Khan --- CHANGELOG.md | 1 + libs/task-commons/build.gradle | 25 ++ .../task/commons/clients/TaskListRequest.java | 103 +++++ .../commons/clients/TaskManagerClient.java | 69 ++++ .../commons/clients/TaskProducerClient.java | 24 ++ .../commons/clients/TaskWorkerClient.java | 38 ++ .../task/commons/clients/package-info.java | 12 + .../opensearch/task/commons/package-info.java | 12 + .../opensearch/task/commons/task/Task.java | 361 ++++++++++++++++++ .../opensearch/task/commons/task/TaskId.java | 55 +++ .../task/commons/task/TaskParams.java | 23 ++ .../task/commons/task/TaskStatus.java | 49 +++ .../task/commons/task/TaskType.java | 27 ++ .../task/commons/task/package-info.java | 12 + .../task/commons/worker/TaskWorker.java | 29 ++ .../task/commons/worker/WorkerNode.java | 96 +++++ .../task/commons/worker/package-info.java | 12 + .../task/commons/mocks/MockTaskParams.java | 25 ++ .../task/commons/task/TaskIdTests.java | 53 +++ .../task/commons/task/TaskTests.java | 147 +++++++ .../task/commons/worker/WorkerNodeTests.java | 52 +++ server/build.gradle | 1 + .../opensearch/common/util/FeatureFlags.java | 5 + .../main/java/org/opensearch/node/Node.java | 12 + .../plugins/TaskManagerClientPlugin.java | 27 ++ .../opensearch/plugins/TaskWorkerPlugin.java | 36 ++ 26 files changed, 1306 insertions(+) create mode 100644 libs/task-commons/build.gradle create mode 100644 libs/task-commons/src/main/java/org/opensearch/task/commons/clients/TaskListRequest.java create mode 100644 libs/task-commons/src/main/java/org/opensearch/task/commons/clients/TaskManagerClient.java create mode 100644 libs/task-commons/src/main/java/org/opensearch/task/commons/clients/TaskProducerClient.java create mode 100644 libs/task-commons/src/main/java/org/opensearch/task/commons/clients/TaskWorkerClient.java create mode 100644 libs/task-commons/src/main/java/org/opensearch/task/commons/clients/package-info.java create mode 100644 libs/task-commons/src/main/java/org/opensearch/task/commons/package-info.java create mode 100644 libs/task-commons/src/main/java/org/opensearch/task/commons/task/Task.java create mode 100644 libs/task-commons/src/main/java/org/opensearch/task/commons/task/TaskId.java create mode 100644 libs/task-commons/src/main/java/org/opensearch/task/commons/task/TaskParams.java create mode 100644 libs/task-commons/src/main/java/org/opensearch/task/commons/task/TaskStatus.java create mode 100644 libs/task-commons/src/main/java/org/opensearch/task/commons/task/TaskType.java create mode 100644 libs/task-commons/src/main/java/org/opensearch/task/commons/task/package-info.java create mode 100644 libs/task-commons/src/main/java/org/opensearch/task/commons/worker/TaskWorker.java create mode 100644 libs/task-commons/src/main/java/org/opensearch/task/commons/worker/WorkerNode.java create mode 100644 libs/task-commons/src/main/java/org/opensearch/task/commons/worker/package-info.java create mode 100644 libs/task-commons/src/test/java/org/opensearch/task/commons/mocks/MockTaskParams.java create mode 100644 libs/task-commons/src/test/java/org/opensearch/task/commons/task/TaskIdTests.java create mode 100644 libs/task-commons/src/test/java/org/opensearch/task/commons/task/TaskTests.java create mode 100644 libs/task-commons/src/test/java/org/opensearch/task/commons/worker/WorkerNodeTests.java create mode 100644 server/src/main/java/org/opensearch/plugins/TaskManagerClientPlugin.java create mode 100644 server/src/main/java/org/opensearch/plugins/TaskWorkerPlugin.java diff --git a/CHANGELOG.md b/CHANGELOG.md index cd0af3932371d..1def0448ad24c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), ## [Unreleased 2.x] ### Added +- [Offline Nodes] Adds offline-tasks library containing various interfaces to be used for Offline Background Tasks. ([#13574](https://github.com/opensearch-project/OpenSearch/pull/13574)) - Fix for hasInitiatedFetching to fix allocation explain and manual reroute APIs (([#14972](https://github.com/opensearch-project/OpenSearch/pull/14972)) - [Workload Management] Add queryGroupId to Task ([14708](https://github.com/opensearch-project/OpenSearch/pull/14708)) - Add setting to ignore throttling nodes for allocation of unassigned primaries in remote restore ([#14991](https://github.com/opensearch-project/OpenSearch/pull/14991)) diff --git a/libs/task-commons/build.gradle b/libs/task-commons/build.gradle new file mode 100644 index 0000000000000..dde3acd8effcf --- /dev/null +++ b/libs/task-commons/build.gradle @@ -0,0 +1,25 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + * + * Modifications Copyright OpenSearch Contributors. See + * GitHub history for details. + */ + +dependencies { + api project(':libs:opensearch-common') + + testImplementation "com.carrotsearch.randomizedtesting:randomizedtesting-runner:${versions.randomizedrunner}" + testImplementation "junit:junit:${versions.junit}" + testImplementation "org.hamcrest:hamcrest:${versions.hamcrest}" + testImplementation(project(":test:framework")) { + exclude group: 'org.opensearch', module: 'opensearch-task-commons' + } +} + +tasks.named('forbiddenApisMain').configure { + replaceSignatureFiles 'jdk-signatures' +} diff --git a/libs/task-commons/src/main/java/org/opensearch/task/commons/clients/TaskListRequest.java b/libs/task-commons/src/main/java/org/opensearch/task/commons/clients/TaskListRequest.java new file mode 100644 index 0000000000000..625e15dfb3b6d --- /dev/null +++ b/libs/task-commons/src/main/java/org/opensearch/task/commons/clients/TaskListRequest.java @@ -0,0 +1,103 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.task.commons.clients; + +import org.opensearch.task.commons.task.TaskStatus; +import org.opensearch.task.commons.task.TaskType; +import org.opensearch.task.commons.worker.WorkerNode; + +/** + * Request object for listing tasks + */ +public class TaskListRequest { + + /** + * Filters listTasks response by specific task status' + */ + private TaskStatus[] taskStatus; + + /** + * Filter listTasks response by specific task types + */ + private TaskType[] taskTypes; + + /** + * Filter listTasks response by specific worker node + */ + private WorkerNode workerNodes; + + /** + * Depicts the start page number for the list call. + * + * @see TaskManagerClient#listTasks(TaskListRequest) + */ + private int startPageNumber; + + /** + * Depicts the page size for the list call. + * + * @see TaskManagerClient#listTasks(TaskListRequest) + */ + private int pageSize; + + /** + * Default constructor + */ + public TaskListRequest() {} + + /** + * Update task types to filter with in the request + * @param taskTypes TaskType[] + * @return ListTaskRequest + */ + public TaskListRequest taskType(TaskType... taskTypes) { + this.taskTypes = taskTypes; + return this; + } + + /** + * Update task status to filter with in the request + * @param taskStatus TaskStatus[] + * @return ListTaskRequest + */ + public TaskListRequest taskType(TaskStatus... taskStatus) { + this.taskStatus = taskStatus; + return this; + } + + /** + * Update worker node to filter with in the request + * @param workerNode WorkerNode + * @return ListTaskRequest + */ + private TaskListRequest workerNode(WorkerNode workerNode) { + this.workerNodes = workerNode; + return this; + } + + /** + * Update page number to start with when fetching the list of tasks + * @param startPageNumber startPageNumber + * @return ListTaskRequest + */ + public TaskListRequest startPageNumber(int startPageNumber) { + this.startPageNumber = startPageNumber; + return this; + } + + /** + * Update page size for the list tasks response + * @param pageSize int + * @return ListTaskRequest + */ + public TaskListRequest pageSize(int pageSize) { + this.pageSize = pageSize; + return this; + } +} diff --git a/libs/task-commons/src/main/java/org/opensearch/task/commons/clients/TaskManagerClient.java b/libs/task-commons/src/main/java/org/opensearch/task/commons/clients/TaskManagerClient.java new file mode 100644 index 0000000000000..23ad8eabdc365 --- /dev/null +++ b/libs/task-commons/src/main/java/org/opensearch/task/commons/clients/TaskManagerClient.java @@ -0,0 +1,69 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.task.commons.clients; + +import org.opensearch.common.annotation.ExperimentalApi; +import org.opensearch.task.commons.task.Task; +import org.opensearch.task.commons.task.TaskId; +import org.opensearch.task.commons.worker.WorkerNode; + +import java.util.List; + +/** + * Client used to interact with Task Store/Queue. + * + * TODO: TaskManager can be something not running an opensearch process. + * We need to come up with a way to allow this interface to be used with in and out opensearch as well + * + * @opensearch.experimental + */ +@ExperimentalApi +public interface TaskManagerClient { + + /** + * Get task from TaskStore/Queue + * + * @param taskId TaskId of the task to be retrieved + * @return Task corresponding to TaskId + */ + Task getTask(TaskId taskId); + + /** + * Update task in TaskStore/Queue + * + * @param task Task to be updated + */ + void updateTask(Task task); + + /** + * Mark task as cancelled. + * Ongoing Tasks can be cancelled as well if the corresponding worker supports cancellation + * + * @param taskId TaskId of the task to be cancelled + */ + void cancelTask(TaskId taskId); + + /** + * List all tasks applying all the filters present in listTaskRequest + * + * @param taskListRequest TaskListRequest + * @return list of all the task matching the filters in listTaskRequest + */ + List listTasks(TaskListRequest taskListRequest); + + /** + * Assign Task to a particular WorkerNode. This ensures no 2 worker Nodes work on the same task. + * This API can be used in both pull and push models of task assignment. + * + * @param taskId TaskId of the task to be assigned + * @param node WorkerNode task is being assigned to + * @return true if task is assigned successfully, false otherwise + */ + boolean assignTask(TaskId taskId, WorkerNode node); +} diff --git a/libs/task-commons/src/main/java/org/opensearch/task/commons/clients/TaskProducerClient.java b/libs/task-commons/src/main/java/org/opensearch/task/commons/clients/TaskProducerClient.java new file mode 100644 index 0000000000000..18f3421f9aa97 --- /dev/null +++ b/libs/task-commons/src/main/java/org/opensearch/task/commons/clients/TaskProducerClient.java @@ -0,0 +1,24 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.task.commons.clients; + +import org.opensearch.task.commons.task.Task; + +/** + * Producer interface used to submit new tasks for execution on worker nodes. + */ +public interface TaskProducerClient { + + /** + * Submit a new task to TaskStore/Queue + * + * @param task Task to be submitted for execution on offline nodes + */ + void submitTask(Task task); +} diff --git a/libs/task-commons/src/main/java/org/opensearch/task/commons/clients/TaskWorkerClient.java b/libs/task-commons/src/main/java/org/opensearch/task/commons/clients/TaskWorkerClient.java new file mode 100644 index 0000000000000..59abe189434dd --- /dev/null +++ b/libs/task-commons/src/main/java/org/opensearch/task/commons/clients/TaskWorkerClient.java @@ -0,0 +1,38 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.task.commons.clients; + +import org.opensearch.task.commons.task.Task; +import org.opensearch.task.commons.task.TaskId; + +import java.util.List; + +/** + * Consumer interface used to find new tasks assigned to a {@code WorkerNode} for execution. + */ +public interface TaskWorkerClient { + + /** + * List all tasks assigned to a WorkerNode. + * Useful when the implementation uses a separate store for Task assignments to Worker nodes + * + * @param taskListRequest TaskListRequest + * @return list of all tasks assigned to a WorkerNode + */ + List getAssignedTasks(TaskListRequest taskListRequest); + + /** + * Sends task heart beat to Task Store/Queue + * + * @param taskId TaskId of Task to send heartbeat for + * @param timestamp timestamp of heartbeat to be recorded in TaskStore/Queue + */ + void sendTaskHeartbeat(TaskId taskId, long timestamp); + +} diff --git a/libs/task-commons/src/main/java/org/opensearch/task/commons/clients/package-info.java b/libs/task-commons/src/main/java/org/opensearch/task/commons/clients/package-info.java new file mode 100644 index 0000000000000..1329f3888248c --- /dev/null +++ b/libs/task-commons/src/main/java/org/opensearch/task/commons/clients/package-info.java @@ -0,0 +1,12 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +/** + * Contains task client related classes + */ +package org.opensearch.task.commons.clients; diff --git a/libs/task-commons/src/main/java/org/opensearch/task/commons/package-info.java b/libs/task-commons/src/main/java/org/opensearch/task/commons/package-info.java new file mode 100644 index 0000000000000..4cb773ace62ce --- /dev/null +++ b/libs/task-commons/src/main/java/org/opensearch/task/commons/package-info.java @@ -0,0 +1,12 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +/** + * Contains offline tasks related classes + */ +package org.opensearch.task.commons; diff --git a/libs/task-commons/src/main/java/org/opensearch/task/commons/task/Task.java b/libs/task-commons/src/main/java/org/opensearch/task/commons/task/Task.java new file mode 100644 index 0000000000000..7ad567b57bd42 --- /dev/null +++ b/libs/task-commons/src/main/java/org/opensearch/task/commons/task/Task.java @@ -0,0 +1,361 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.task.commons.task; + +import org.opensearch.common.Nullable; +import org.opensearch.common.annotation.ExperimentalApi; +import org.opensearch.task.commons.worker.WorkerNode; + +/** + * A Background Task to be run on Offline Node. + */ +@ExperimentalApi +public class Task { + + /** + * Task identifier used to uniquely identify a Task + */ + private final TaskId taskId; + + /** + * Depicts latest state of the Task + */ + private final TaskStatus taskStatus; + + /** + * Various params to used for Task execution + */ + private final TaskParams params; + + /** + * Type/Category of the Task + */ + private final TaskType taskType; + + /** + * Worker Node on which the Task is to be executed + */ + private final WorkerNode assignedNode; + + /** + * Timestamp at which the Task was created + */ + private final long createdAt; + + /** + * Timestamp at which the Task was assigned to a worker + */ + private final long assignedAt; + + /** + * Timestamp at which the Task was started execution on worker + */ + private final long startedAt; + + /** + * Timestamp at which the Task was either completed/failed/cancelled + */ + private final long completedAt; + + /** + * Timestamp at which last heartbeat was sent by the worker + */ + private final long lastHeartbeatAt; + + /** + * Constructor for Task + * + * @param taskId Task identifier + * @param taskStatus Task status + * @param params Task Params + * @param taskType Task Type + * @param createdAt Timestamp at which the Task was created + * @param assignedAt Timestamp at which the Task was assigned to a worker + * @param startedAt Timestamp at which the Task was started execution on worker + * @param completedAt Timestamp at which the Task was either completed/failed/cancelled + * @param lastHeartbeatAt Timestamp at which last heartbeat was sent by the worker + * @param assignedNode Worker Node on which the Task is to be executed + */ + public Task( + TaskId taskId, + TaskStatus taskStatus, + TaskParams params, + TaskType taskType, + long createdAt, + @Nullable long assignedAt, + @Nullable long startedAt, + @Nullable long completedAt, + @Nullable long lastHeartbeatAt, + @Nullable WorkerNode assignedNode + ) { + this.taskId = taskId; + this.taskStatus = taskStatus; + this.params = params; + this.taskType = taskType; + this.createdAt = createdAt; + this.assignedAt = assignedAt; + this.startedAt = startedAt; + this.completedAt = completedAt; + this.lastHeartbeatAt = lastHeartbeatAt; + this.assignedNode = assignedNode; + } + + /** + * Get TaskId + * @return TaskId + */ + public TaskId getTaskId() { + return taskId; + } + + /** + * Get TaskStatus + * @return TaskStatus + */ + public TaskStatus getTaskStatus() { + return taskStatus; + } + + /** + * Get TaskParams + * @return TaskParams + */ + public TaskParams getParams() { + return params; + } + + /** + * Get TaskType + * @return TaskType + */ + public TaskType getTaskType() { + return taskType; + } + + /** + * Get Task Creation Time + * @return createdAt + */ + public long getCreatedAt() { + return createdAt; + } + + /** + * Get Task Assignment Time + * @return assignedAt + */ + public long getAssignedAt() { + return assignedAt; + } + + /** + * Get Task Start Time + * @return startedAt + */ + public long getStartedAt() { + return startedAt; + } + + /** + * Get Task Completion Time + * @return completedAt + */ + public long getCompletedAt() { + return completedAt; + } + + /** + * Get Last Heartbeat Time + * @return lastHeartbeatAt + */ + public long getLastHeartbeatAt() { + return lastHeartbeatAt; + } + + /** + * Get Task Assigned Node + * @return assignedNode + */ + public WorkerNode getAssignedNode() { + return assignedNode; + } + + /** + * Builder class for Task. + */ + public static class Builder { + /** + * Task identifier used to uniquely identify a Task + */ + private final TaskId taskId; + + /** + * Depicts latest state of the Task + */ + private TaskStatus taskStatus; + + /** + * Various params to used for Task execution + */ + private final TaskParams params; + + /** + * Type/Category of the Task + */ + private final TaskType taskType; + + /** + * Type/Category of the Task + */ + private WorkerNode assignedNode; + + /** + * Timestamp at which the Task was created + */ + private final long createdAt; + + /** + * Timestamp at which the Task was assigned to a worker + */ + private long assignedAt; + + /** + * Timestamp at which the Task was started execution on worker + */ + private long startedAt; + + /** + * Timestamp at which the Task was either completed/failed/cancelled + */ + private long completedAt; + + /** + * Timestamp at which last heartbeat was sent by the worker + */ + private long lastHeartbeatAt; + + /** + * Constructor for Task Builder + * + * @param taskId Task identifier + * @param taskStatus Task status + * @param params Task Params + * @param taskType Task Type + * @param createdAt Task Creation Time + */ + private Builder(TaskId taskId, TaskStatus taskStatus, TaskParams params, TaskType taskType, long createdAt) { + this.taskId = taskId; + this.taskStatus = taskStatus; + this.params = params; + this.taskType = taskType; + this.createdAt = createdAt; + } + + /** + * Build Builder from Task + * @param task Task to build from + * @return Task.Builder + */ + public static Builder builder(Task task) { + Builder builder = new Builder( + task.getTaskId(), + task.getTaskStatus(), + task.getParams(), + task.getTaskType(), + task.getCreatedAt() + ); + builder.assignedAt(task.getAssignedAt()); + builder.startedAt(task.getStartedAt()); + builder.completedAt(task.getCompletedAt()); + builder.lastHeartbeatAt(task.getLastHeartbeatAt()); + builder.assignedNode(task.getAssignedNode()); + return builder; + } + + /** + * Build Builder from various Task attributes + * @param taskId Task identifier + * @param taskStatus TaskStatus + * @param params TaskParams + * @param taskType TaskType + * @param createdAt Task Creation Time + * @return Task.Builder + */ + public static Builder builder(TaskId taskId, TaskStatus taskStatus, TaskParams params, TaskType taskType, long createdAt) { + return new Builder(taskId, taskStatus, params, taskType, createdAt); + } + + /** + * Set Task Assignment Time + * @param assignedAt Timestamp at which the Task was assigned to a worker + */ + public void assignedAt(long assignedAt) { + this.assignedAt = assignedAt; + } + + /** + * Set Task Start Time + * @param startedAt Timestamp at which the Task was started execution on worker + */ + public void startedAt(long startedAt) { + this.startedAt = startedAt; + } + + /** + * Set Task Completion Time + * @param completedAt Timestamp at which the Task was either completed/failed/cancelled + */ + public void completedAt(long completedAt) { + this.completedAt = completedAt; + } + + /** + * Set Task Last Heartbeat Time for the task + * @param lastHeartbeatAt Timestamp at which last heartbeat was sent by the worker + */ + public void lastHeartbeatAt(long lastHeartbeatAt) { + this.lastHeartbeatAt = lastHeartbeatAt; + } + + /** + * Update the Task Status + * @param taskStatus {@link TaskStatus} - current status of the Task + */ + public void taskStatus(TaskStatus taskStatus) { + this.taskStatus = taskStatus; + } + + /** + * Set Task Assigned Node + * @param node Worker Node on which the Task is to be executed + */ + public void assignedNode(WorkerNode node) { + this.assignedNode = node; + } + + /** + * Build Task from Builder + * @return Task + */ + public Task build() { + return new Task( + taskId, + taskStatus, + params, + taskType, + createdAt, + assignedAt, + startedAt, + completedAt, + lastHeartbeatAt, + assignedNode + ); + } + } +} diff --git a/libs/task-commons/src/main/java/org/opensearch/task/commons/task/TaskId.java b/libs/task-commons/src/main/java/org/opensearch/task/commons/task/TaskId.java new file mode 100644 index 0000000000000..7fb7e1536dc17 --- /dev/null +++ b/libs/task-commons/src/main/java/org/opensearch/task/commons/task/TaskId.java @@ -0,0 +1,55 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.task.commons.task; + +import org.opensearch.common.annotation.ExperimentalApi; + +import java.util.Objects; + +/** + * Class encapsulating Task identifier + */ +@ExperimentalApi +public class TaskId { + + /** + * Identified of the Task + */ + private final String id; + + /** + * Constructor to initialize TaskId + * @param id String value of Task id + */ + public TaskId(String id) { + this.id = id; + } + + /** + * Get id value + * @return id + */ + public String getValue() { + return id; + } + + @Override + public boolean equals(Object obj) { + if (obj == null || getClass() != obj.getClass()) { + return false; + } + TaskId other = (TaskId) obj; + return this.id.equals(other.id); + } + + @Override + public int hashCode() { + return Objects.hash(id); + } +} diff --git a/libs/task-commons/src/main/java/org/opensearch/task/commons/task/TaskParams.java b/libs/task-commons/src/main/java/org/opensearch/task/commons/task/TaskParams.java new file mode 100644 index 0000000000000..ac6126a02cf24 --- /dev/null +++ b/libs/task-commons/src/main/java/org/opensearch/task/commons/task/TaskParams.java @@ -0,0 +1,23 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.task.commons.task; + +import org.opensearch.common.annotation.ExperimentalApi; + +/** + * Base class for all TaskParams implementation of various TaskTypes + */ +@ExperimentalApi +public abstract class TaskParams { + + /** + * Default constructor + */ + public TaskParams() {} +} diff --git a/libs/task-commons/src/main/java/org/opensearch/task/commons/task/TaskStatus.java b/libs/task-commons/src/main/java/org/opensearch/task/commons/task/TaskStatus.java new file mode 100644 index 0000000000000..e44e5c7fee620 --- /dev/null +++ b/libs/task-commons/src/main/java/org/opensearch/task/commons/task/TaskStatus.java @@ -0,0 +1,49 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.task.commons.task; + +import org.opensearch.common.annotation.ExperimentalApi; + +/** + * Task status enum + */ +@ExperimentalApi +public enum TaskStatus { + + /** + * TaskStatus of a Task which is not yet assigned to or picked up by a worker + */ + UNASSIGNED, + + /** + * TaskStatus of a Task which is assigned to or picked up by a worker but hasn't started execution yet. + * This status confirms that a worker will execute this task and no other worker should pick it up. + */ + ASSIGNED, + + /** + * TaskStatus of an in progress Task + */ + ACTIVE, + + /** + * TaskStatus of a finished Task + */ + SUCCESS, + + /** + * TaskStatus of a Task which failed in 1 or more attempts + */ + FAILED, + + /** + * TaskStatus of a cancelled Task + */ + CANCELLED +} diff --git a/libs/task-commons/src/main/java/org/opensearch/task/commons/task/TaskType.java b/libs/task-commons/src/main/java/org/opensearch/task/commons/task/TaskType.java new file mode 100644 index 0000000000000..ef267a86e3da4 --- /dev/null +++ b/libs/task-commons/src/main/java/org/opensearch/task/commons/task/TaskType.java @@ -0,0 +1,27 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.task.commons.task; + +import org.opensearch.common.annotation.ExperimentalApi; + +/** + * Enum for task type + */ +@ExperimentalApi +public enum TaskType { + /** + * For all segment merge related tasks + */ + MERGE, + + /** + * For all snapshot related tasks + */ + SNAPSHOT +} diff --git a/libs/task-commons/src/main/java/org/opensearch/task/commons/task/package-info.java b/libs/task-commons/src/main/java/org/opensearch/task/commons/task/package-info.java new file mode 100644 index 0000000000000..e99e4079cc7af --- /dev/null +++ b/libs/task-commons/src/main/java/org/opensearch/task/commons/task/package-info.java @@ -0,0 +1,12 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +/** + * Contains tasks related classes + */ +package org.opensearch.task.commons.task; diff --git a/libs/task-commons/src/main/java/org/opensearch/task/commons/worker/TaskWorker.java b/libs/task-commons/src/main/java/org/opensearch/task/commons/worker/TaskWorker.java new file mode 100644 index 0000000000000..6d9fd1fed23e4 --- /dev/null +++ b/libs/task-commons/src/main/java/org/opensearch/task/commons/worker/TaskWorker.java @@ -0,0 +1,29 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.task.commons.worker; + +import org.opensearch.common.annotation.ExperimentalApi; +import org.opensearch.task.commons.task.Task; + +/** + * Task Worker that executes the Task + * + * @opensearch.experimental + */ +@ExperimentalApi +public interface TaskWorker { + + /** + * Execute the Task + * + * @param task Task to be executed + */ + void executeTask(Task task); + +} diff --git a/libs/task-commons/src/main/java/org/opensearch/task/commons/worker/WorkerNode.java b/libs/task-commons/src/main/java/org/opensearch/task/commons/worker/WorkerNode.java new file mode 100644 index 0000000000000..1b9426d472184 --- /dev/null +++ b/libs/task-commons/src/main/java/org/opensearch/task/commons/worker/WorkerNode.java @@ -0,0 +1,96 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.task.commons.worker; + +import java.util.Objects; + +/** + * Represents a worker node in the fleet + */ +public class WorkerNode { + /** + * The unique identifier of the worker node. + */ + private final String id; + + /** + * The name of the worker node. + */ + private final String name; + + /** + * The IP address of the worker node. + */ + private final String ip; + + /** + * Creates a new worker node with the given ID, name, and IP address. + * + * @param id The unique identifier of the worker node. + * @param name The name of the worker node. + * @param ip The IP address of the worker node. + */ + private WorkerNode(String id, String name, String ip) { + this.id = id; + this.name = name; + this.ip = ip; + } + + /** + * Creates a new worker node with the given ID, name, and IP address. + * + * @param id The unique identifier of the worker node. + * @param name The name of the worker node. + * @param ip The IP address of the worker node. + * @return The created worker node. + */ + public static WorkerNode createWorkerNode(String id, String name, String ip) { + return new WorkerNode(id, name, ip); + } + + /** + * Returns the unique identifier of the worker node. + * + * @return The ID of the worker node. + */ + public String getId() { + return id; + } + + /** + * Returns the name of the worker node. + * + * @return The name of the worker node. + */ + public String getName() { + return name; + } + + /** + * Returns the IP address of the worker node. + * + * @return The IP address of the worker node. + */ + public String getIp() { + return ip; + } + + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + WorkerNode that = (WorkerNode) o; + return Objects.equals(id, that.id) && Objects.equals(name, that.name) && Objects.equals(ip, that.ip); + } + + @Override + public int hashCode() { + return Objects.hash(id, name, ip); + } +} diff --git a/libs/task-commons/src/main/java/org/opensearch/task/commons/worker/package-info.java b/libs/task-commons/src/main/java/org/opensearch/task/commons/worker/package-info.java new file mode 100644 index 0000000000000..d74fa30e8b661 --- /dev/null +++ b/libs/task-commons/src/main/java/org/opensearch/task/commons/worker/package-info.java @@ -0,0 +1,12 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +/** + * Contains task worker related classes + */ +package org.opensearch.task.commons.worker; diff --git a/libs/task-commons/src/test/java/org/opensearch/task/commons/mocks/MockTaskParams.java b/libs/task-commons/src/test/java/org/opensearch/task/commons/mocks/MockTaskParams.java new file mode 100644 index 0000000000000..21d2a3dd725b6 --- /dev/null +++ b/libs/task-commons/src/test/java/org/opensearch/task/commons/mocks/MockTaskParams.java @@ -0,0 +1,25 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.task.commons.mocks; + +import org.opensearch.task.commons.task.TaskParams; + +public class MockTaskParams extends TaskParams { + + private final String value; + + public MockTaskParams(String mockValue) { + super(); + value = mockValue; + } + + public String getValue() { + return value; + } +} diff --git a/libs/task-commons/src/test/java/org/opensearch/task/commons/task/TaskIdTests.java b/libs/task-commons/src/test/java/org/opensearch/task/commons/task/TaskIdTests.java new file mode 100644 index 0000000000000..e58382e9fc056 --- /dev/null +++ b/libs/task-commons/src/test/java/org/opensearch/task/commons/task/TaskIdTests.java @@ -0,0 +1,53 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.task.commons.task; + +import org.opensearch.test.OpenSearchTestCase; + +/** + * Tests for {@link TaskId} + */ +public class TaskIdTests extends OpenSearchTestCase { + + public void testConstructorAndGetValue() { + TaskId taskId = new TaskId("123"); + assertEquals("123", taskId.getValue()); + } + + public void testEqualsWithSameId() { + TaskId taskId1 = new TaskId("456"); + TaskId taskId2 = new TaskId("456"); + assertEquals(taskId1, taskId2); + } + + public void testEqualsWithDifferentId() { + TaskId taskId1 = new TaskId("789"); + TaskId taskId2 = new TaskId("987"); + assertNotEquals(taskId1, taskId2); + } + + public void testEqualsWithNull() { + TaskId taskId = new TaskId("abc"); + assertNotEquals(null, taskId); + } + + public void testEqualsWithDifferentClass() { + TaskId taskId = new TaskId("def"); + assertNotEquals(taskId, new Object()); + } + + public void testHashCode() { + TaskId taskId1 = new TaskId("456"); + TaskId taskId2 = new TaskId("456"); + assertEquals(taskId1.hashCode(), taskId2.hashCode()); + + TaskId taskId3 = new TaskId("4567"); + assertNotEquals(taskId1.hashCode(), taskId3.hashCode()); + } +} diff --git a/libs/task-commons/src/test/java/org/opensearch/task/commons/task/TaskTests.java b/libs/task-commons/src/test/java/org/opensearch/task/commons/task/TaskTests.java new file mode 100644 index 0000000000000..37c5f543daa02 --- /dev/null +++ b/libs/task-commons/src/test/java/org/opensearch/task/commons/task/TaskTests.java @@ -0,0 +1,147 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.task.commons.task; + +import org.opensearch.task.commons.mocks.MockTaskParams; +import org.opensearch.task.commons.worker.WorkerNode; +import org.opensearch.test.OpenSearchTestCase; + +/** + * Test for {@link Task} + */ +public class TaskTests extends OpenSearchTestCase { + + public void testTaskConstructorAndGetters() { + TaskId taskId = new TaskId("123"); + TaskStatus taskStatus = TaskStatus.UNASSIGNED; + TaskParams params = new MockTaskParams("mock"); + TaskType taskType = TaskType.MERGE; + long createdAt = System.currentTimeMillis(); + long assignedAt = createdAt + 1000; + long startedAt = createdAt + 2000; + long completedAt = createdAt + 3000; + long lastHeartbeatAt = createdAt + 2500; + WorkerNode assignedNode = WorkerNode.createWorkerNode("node1", "nodeip", "nodename"); + + Task task = new Task( + taskId, + taskStatus, + params, + taskType, + createdAt, + assignedAt, + startedAt, + completedAt, + lastHeartbeatAt, + assignedNode + ); + + assertEquals(taskId, task.getTaskId()); + assertEquals(taskStatus, task.getTaskStatus()); + assertEquals(params, task.getParams()); + assertEquals(taskType, task.getTaskType()); + assertEquals(createdAt, task.getCreatedAt()); + assertEquals(assignedAt, task.getAssignedAt()); + assertEquals(startedAt, task.getStartedAt()); + assertEquals(completedAt, task.getCompletedAt()); + assertEquals(lastHeartbeatAt, task.getLastHeartbeatAt()); + assertEquals(assignedNode, task.getAssignedNode()); + } + + public void testBuilderFromTask() { + TaskId taskId = new TaskId("123"); + TaskStatus taskStatus = TaskStatus.UNASSIGNED; + TaskParams params = new MockTaskParams("mock"); + TaskType taskType = TaskType.MERGE; + long createdAt = System.currentTimeMillis(); + long assignedAt = createdAt + 1000; + long startedAt = createdAt + 2000; + long completedAt = createdAt + 3000; + long lastHeartbeatAt = createdAt + 2500; + WorkerNode assignedNode = WorkerNode.createWorkerNode("node1", "nodeip", "nodename"); + + Task originalTask = new Task( + taskId, + taskStatus, + params, + taskType, + createdAt, + assignedAt, + startedAt, + completedAt, + lastHeartbeatAt, + assignedNode + ); + + Task.Builder builder = Task.Builder.builder(originalTask); + Task newTask = builder.build(); + + assertEquals(originalTask.getTaskId(), newTask.getTaskId()); + assertEquals(originalTask.getTaskStatus(), newTask.getTaskStatus()); + assertEquals(originalTask.getParams(), newTask.getParams()); + assertEquals(originalTask.getTaskType(), newTask.getTaskType()); + assertEquals(originalTask.getCreatedAt(), newTask.getCreatedAt()); + assertEquals(originalTask.getAssignedAt(), newTask.getAssignedAt()); + assertEquals(originalTask.getStartedAt(), newTask.getStartedAt()); + assertEquals(originalTask.getCompletedAt(), newTask.getCompletedAt()); + assertEquals(originalTask.getLastHeartbeatAt(), newTask.getLastHeartbeatAt()); + assertEquals(originalTask.getAssignedNode(), newTask.getAssignedNode()); + } + + public void testBuilderFromAttributes() { + TaskId taskId = new TaskId("123"); + TaskStatus taskStatus = TaskStatus.UNASSIGNED; + TaskParams params = new MockTaskParams("mock"); + TaskType taskType = TaskType.MERGE; + long createdAt = System.currentTimeMillis(); + + Task.Builder builder = Task.Builder.builder(taskId, taskStatus, params, taskType, createdAt); + builder.assignedAt(createdAt + 1000); + builder.startedAt(createdAt + 2000); + builder.completedAt(createdAt + 3000); + builder.lastHeartbeatAt(createdAt + 2500); + builder.assignedNode(WorkerNode.createWorkerNode("node1", "nodeip", "nodename")); + builder.taskStatus(TaskStatus.ACTIVE); + + Task task = builder.build(); + + assertEquals(taskId, task.getTaskId()); + assertEquals(TaskStatus.ACTIVE, task.getTaskStatus()); + assertEquals(params, task.getParams()); + assertEquals(taskType, task.getTaskType()); + assertEquals(createdAt, task.getCreatedAt()); + assertEquals(createdAt + 1000, task.getAssignedAt()); + assertEquals(createdAt + 2000, task.getStartedAt()); + assertEquals(createdAt + 3000, task.getCompletedAt()); + assertEquals(createdAt + 2500, task.getLastHeartbeatAt()); + assertEquals(WorkerNode.createWorkerNode("node1", "nodeip", "nodename"), task.getAssignedNode()); + } + + public void testBuilderWithNullOptionalFields() { + TaskId taskId = new TaskId("123"); + TaskStatus taskStatus = TaskStatus.UNASSIGNED; + TaskParams params = new MockTaskParams("mock"); + TaskType taskType = TaskType.MERGE; + long createdAt = System.currentTimeMillis(); + + Task.Builder builder = Task.Builder.builder(taskId, taskStatus, params, taskType, createdAt); + Task task = builder.build(); + + assertEquals(taskId, task.getTaskId()); + assertEquals(taskStatus, task.getTaskStatus()); + assertEquals(params, task.getParams()); + assertEquals(taskType, task.getTaskType()); + assertEquals(createdAt, task.getCreatedAt()); + assertEquals(0, task.getAssignedAt()); + assertEquals(0, task.getStartedAt()); + assertEquals(0, task.getCompletedAt()); + assertEquals(0, task.getLastHeartbeatAt()); + assertNull(task.getAssignedNode()); + } +} diff --git a/libs/task-commons/src/test/java/org/opensearch/task/commons/worker/WorkerNodeTests.java b/libs/task-commons/src/test/java/org/opensearch/task/commons/worker/WorkerNodeTests.java new file mode 100644 index 0000000000000..07570bf947a27 --- /dev/null +++ b/libs/task-commons/src/test/java/org/opensearch/task/commons/worker/WorkerNodeTests.java @@ -0,0 +1,52 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.task.commons.worker; + +import org.opensearch.test.OpenSearchTestCase; + +/** + * Tests for {@link WorkerNode} + */ +public class WorkerNodeTests extends OpenSearchTestCase { + + public void testCreateWorkerNode() { + WorkerNode worker = WorkerNode.createWorkerNode("1", "Worker1", "192.168.1.1"); + assertNotNull(worker); + assertEquals("1", worker.getId()); + assertEquals("Worker1", worker.getName()); + assertEquals("192.168.1.1", worker.getIp()); + } + + public void testEquality() { + WorkerNode worker1 = WorkerNode.createWorkerNode("5", "Worker5", "192.168.1.5"); + WorkerNode worker2 = WorkerNode.createWorkerNode("5", "Worker5", "192.168.1.5"); + WorkerNode worker3 = WorkerNode.createWorkerNode("6", "Worker6", "192.168.1.6"); + + assertEquals(worker1, worker2); + assertNotEquals(worker1, worker3); + assertNotEquals(worker2, worker3); + } + + public void testHashCode() { + WorkerNode worker1 = WorkerNode.createWorkerNode("7", "Worker7", "192.168.1.7"); + WorkerNode worker2 = WorkerNode.createWorkerNode("7", "Worker7", "192.168.1.7"); + + assertEquals(worker1.hashCode(), worker2.hashCode()); + } + + public void testNotEqualToNull() { + WorkerNode worker = WorkerNode.createWorkerNode("8", "Worker8", "192.168.1.8"); + assertNotEquals(null, worker); + } + + public void testNotEqualToDifferentClass() { + WorkerNode worker = WorkerNode.createWorkerNode("9", "Worker9", "192.168.1.9"); + assertNotEquals(worker, new Object()); + } +} diff --git a/server/build.gradle b/server/build.gradle index 429af5d0ac258..5facc73dff968 100644 --- a/server/build.gradle +++ b/server/build.gradle @@ -68,6 +68,7 @@ dependencies { api project(':libs:opensearch-x-content') api project(":libs:opensearch-geo") api project(":libs:opensearch-telemetry") + api project(":libs:opensearch-task-commons") compileOnly project(':libs:opensearch-plugin-classloader') diff --git a/server/src/main/java/org/opensearch/common/util/FeatureFlags.java b/server/src/main/java/org/opensearch/common/util/FeatureFlags.java index 9d57e6939e3ae..e2554d61116ad 100644 --- a/server/src/main/java/org/opensearch/common/util/FeatureFlags.java +++ b/server/src/main/java/org/opensearch/common/util/FeatureFlags.java @@ -72,6 +72,11 @@ public class FeatureFlags { */ public static final String REMOTE_PUBLICATION_EXPERIMENTAL = "opensearch.experimental.feature.remote_store.publication.enabled"; + /** + * Gates the functionality of background task execution. + */ + public static final String BACKGROUND_TASK_EXECUTION_EXPERIMENTAL = "opensearch.experimental.feature.task.background.enabled"; + public static final Setting REMOTE_STORE_MIGRATION_EXPERIMENTAL_SETTING = Setting.boolSetting( REMOTE_STORE_MIGRATION_EXPERIMENTAL, false, diff --git a/server/src/main/java/org/opensearch/node/Node.java b/server/src/main/java/org/opensearch/node/Node.java index cbed8dfea8cc4..409f84354a8b1 100644 --- a/server/src/main/java/org/opensearch/node/Node.java +++ b/server/src/main/java/org/opensearch/node/Node.java @@ -215,6 +215,7 @@ import org.opensearch.plugins.SearchPlugin; import org.opensearch.plugins.SecureSettingsFactory; import org.opensearch.plugins.SystemIndexPlugin; +import org.opensearch.plugins.TaskManagerClientPlugin; import org.opensearch.plugins.TelemetryAwarePlugin; import org.opensearch.plugins.TelemetryPlugin; import org.opensearch.ratelimitting.admissioncontrol.AdmissionControlService; @@ -239,6 +240,7 @@ import org.opensearch.snapshots.SnapshotShardsService; import org.opensearch.snapshots.SnapshotsInfoService; import org.opensearch.snapshots.SnapshotsService; +import org.opensearch.task.commons.clients.TaskManagerClient; import org.opensearch.tasks.Task; import org.opensearch.tasks.TaskCancellationMonitoringService; import org.opensearch.tasks.TaskCancellationMonitoringSettings; @@ -298,6 +300,7 @@ import java.util.stream.Stream; import static java.util.stream.Collectors.toList; +import static org.opensearch.common.util.FeatureFlags.BACKGROUND_TASK_EXECUTION_EXPERIMENTAL; import static org.opensearch.common.util.FeatureFlags.TELEMETRY; import static org.opensearch.env.NodeEnvironment.collectFileCacheDataPath; import static org.opensearch.index.ShardIndexingPressureSettings.SHARD_INDEXING_PRESSURE_ENABLED_ATTRIBUTE_KEY; @@ -1314,6 +1317,13 @@ protected Node( .flatMap(List::stream) .collect(toList()); + final Optional taskManagerClientOptional = FeatureFlags.isEnabled(BACKGROUND_TASK_EXECUTION_EXPERIMENTAL) + ? pluginsService.filterPlugins(TaskManagerClientPlugin.class) + .stream() + .map(plugin -> plugin.getTaskManagerClient(client, clusterService, threadPool)) + .findFirst() + : Optional.empty(); + final PersistentTasksExecutorRegistry registry = new PersistentTasksExecutorRegistry(tasksExecutors); final PersistentTasksClusterService persistentTasksClusterService = new PersistentTasksClusterService( settings, @@ -1420,6 +1430,8 @@ protected Node( b.bind(PersistedStateRegistry.class).toInstance(persistedStateRegistry); b.bind(SegmentReplicationStatsTracker.class).toInstance(segmentReplicationStatsTracker); b.bind(SearchRequestOperationsCompositeListenerFactory.class).toInstance(searchRequestOperationsCompositeListenerFactory); + + taskManagerClientOptional.ifPresent(value -> b.bind(TaskManagerClient.class).toInstance(value)); }); injector = modules.createInjector(); diff --git a/server/src/main/java/org/opensearch/plugins/TaskManagerClientPlugin.java b/server/src/main/java/org/opensearch/plugins/TaskManagerClientPlugin.java new file mode 100644 index 0000000000000..3f511bd4c48d2 --- /dev/null +++ b/server/src/main/java/org/opensearch/plugins/TaskManagerClientPlugin.java @@ -0,0 +1,27 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.plugins; + +import org.opensearch.client.Client; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.annotation.ExperimentalApi; +import org.opensearch.task.commons.clients.TaskManagerClient; +import org.opensearch.threadpool.ThreadPool; + +/** + * Plugin to provide an implementation of Task client + */ +@ExperimentalApi +public interface TaskManagerClientPlugin { + + /** + * Get the task client. + */ + TaskManagerClient getTaskManagerClient(Client client, ClusterService clusterService, ThreadPool threadPool); +} diff --git a/server/src/main/java/org/opensearch/plugins/TaskWorkerPlugin.java b/server/src/main/java/org/opensearch/plugins/TaskWorkerPlugin.java new file mode 100644 index 0000000000000..3ac96ed989425 --- /dev/null +++ b/server/src/main/java/org/opensearch/plugins/TaskWorkerPlugin.java @@ -0,0 +1,36 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.plugins; + +import org.opensearch.client.Client; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.annotation.ExperimentalApi; +import org.opensearch.task.commons.task.TaskType; +import org.opensearch.task.commons.worker.TaskWorker; +import org.opensearch.threadpool.ThreadPool; + +/** + * Plugin for providing TaskWorkers for Offline Nodes + */ +@ExperimentalApi +public interface TaskWorkerPlugin { + + /** + * Get the new TaskWorker for a TaskType + * + * @return TaskWorker to execute Tasks on Offline Nodes + */ + TaskWorker getTaskWorker(Client client, ClusterService clusterService, ThreadPool threadPool); + + /** + * Get the TaskType for this TaskWorker + * @return TaskType + */ + TaskType getTaskType(); +} From 01acf1c189293784205db7589fb435c3ab1c89fe Mon Sep 17 00:00:00 2001 From: Andriy Redko Date: Wed, 14 Aug 2024 17:43:43 -0400 Subject: [PATCH 34/40] Update to Gradle 8.10 (#15177) * Update to Gradle 8.10 Signed-off-by: Andriy Redko * Address code review comments Signed-off-by: Andriy Redko --------- Signed-off-by: Andriy Redko --- .github/workflows/assemble.yml | 16 ++++++++++++++-- .github/workflows/precommit.yml | 14 +++++++++++++- gradle/wrapper/gradle-wrapper.jar | Bin 43504 -> 43583 bytes gradle/wrapper/gradle-wrapper.properties | 4 ++-- 4 files changed, 29 insertions(+), 5 deletions(-) diff --git a/.github/workflows/assemble.yml b/.github/workflows/assemble.yml index 51ae075ffa2c9..200d9cf2d788b 100644 --- a/.github/workflows/assemble.yml +++ b/.github/workflows/assemble.yml @@ -16,6 +16,17 @@ jobs: with: java-version: ${{ matrix.java }} distribution: temurin + - name: Set up JDK 17 + # See please https://docs.gradle.org/8.10/userguide/upgrading_version_8.html#minimum_daemon_jvm_version + if: matrix.java == 11 + uses: actions/setup-java@v4 + with: + java-version: 17 + distribution: temurin + - name: Set JAVA${{ matrix.java }}_HOME + shell: bash + run: | + echo "JAVA${{ matrix.java }}_HOME=$JAVA_HOME_${{ matrix.java }}_${{ runner.arch }}" >> $GITHUB_ENV - name: Setup docker (missing on MacOS) id: setup_docker if: runner.os == 'macos' @@ -30,10 +41,11 @@ jobs: # Report success even if previous step failed (Docker on MacOS runner is very unstable) exit 0; - name: Run Gradle (assemble) + shell: bash if: runner.os != 'macos' run: | - ./gradlew assemble --parallel --no-build-cache -PDISABLE_BUILD_CACHE + ./gradlew assemble --parallel --no-build-cache -PDISABLE_BUILD_CACHE -Druntime.java=${{ matrix.java }} - name: Run Gradle (assemble) if: runner.os == 'macos' && steps.setup_docker.outcome == 'success' run: | - ./gradlew assemble --parallel --no-build-cache -PDISABLE_BUILD_CACHE + ./gradlew assemble --parallel --no-build-cache -PDISABLE_BUILD_CACHE -Druntime.java=${{ matrix.java }} diff --git a/.github/workflows/precommit.yml b/.github/workflows/precommit.yml index 793fdae5df4da..7c65df1f677a5 100644 --- a/.github/workflows/precommit.yml +++ b/.github/workflows/precommit.yml @@ -17,6 +17,18 @@ jobs: java-version: ${{ matrix.java }} distribution: temurin cache: gradle + - name: Set up JDK 17 + # See please https://docs.gradle.org/8.10/userguide/upgrading_version_8.html#minimum_daemon_jvm_version + if: matrix.java == 11 + uses: actions/setup-java@v4 + with: + java-version: 17 + distribution: temurin + - name: Set JAVA${{ matrix.java }}_HOME + shell: bash + run: | + echo "JAVA${{ matrix.java }}_HOME=$JAVA_HOME_${{ matrix.java }}_${{ runner.arch }}" >> $GITHUB_ENV - name: Run Gradle (precommit) + shell: bash run: | - ./gradlew javadoc precommit --parallel + ./gradlew javadoc precommit --parallel -Druntime.java=${{ matrix.java }} diff --git a/gradle/wrapper/gradle-wrapper.jar b/gradle/wrapper/gradle-wrapper.jar index 2c3521197d7c4586c843d1d3e9090525f1898cde..a4b76b9530d66f5e68d973ea569d8e19de379189 100644 GIT binary patch delta 3990 zcmV;H4{7l5(*nQL0Kr1kzC=_KMxQY0|W5(lc#i zH*M1^P4B}|{x<+fkObwl)u#`$GxKKV&3pg*-y6R6txw)0qU|Clf9Uds3x{_-**c=7 z&*)~RHPM>Rw#Hi1R({;bX|7?J@w}DMF>dQQU2}9yj%iLjJ*KD6IEB2^n#gK7M~}6R zkH+)bc--JU^pV~7W=3{E*4|ZFpDpBa7;wh4_%;?XM-5ZgZNnVJ=vm!%a2CdQb?oTa z70>8rTb~M$5Tp!Se+4_OKWOB1LF+7gv~$$fGC95ToUM(I>vrd$>9|@h=O?eARj0MH zT4zo(M>`LWoYvE>pXvqG=d96D-4?VySz~=tPVNyD$XMshoTX(1ZLB5OU!I2OI{kb) zS8$B8Qm>wLT6diNnyJZC?yp{Kn67S{TCOt-!OonOK7$K)e-13U9GlnQXPAb&SJ0#3 z+vs~+4Qovv(%i8g$I#FCpCG^C4DdyQw3phJ(f#y*pvNDQCRZ~MvW<}fUs~PL=4??j zmhPyg<*I4RbTz|NHFE-DC7lf2=}-sGkE5e!RM%3ohM7_I^IF=?O{m*uUPH(V?gqyc(Rp?-Qu(3bBIL4Fz(v?=_Sh?LbK{nqZMD>#9D_hNhaV$0ef3@9V90|0u#|PUNTO>$F=qRhg1duaE z0`v~X3G{8RVT@kOa-pU+z8{JWyP6GF*u2e8eKr7a2t1fuqQy)@d|Qn(%YLZ62TWtoX@$nL}9?atE#Yw`rd(>cr0gY;dT9~^oL;u)zgHUvxc2I*b&ZkGM-iq=&(?kyO(3}=P! zRp=rErEyMT5UE9GjPHZ#T<`cnD)jyIL!8P{H@IU#`e8cAG5jMK zVyKw7--dAC;?-qEu*rMr$5@y535qZ6p(R#+fLA_)G~!wnT~~)|s`}&fA(s6xXN`9j zP#Fd3GBa#HeS{5&8p?%DKUyN^X9cYUc6vq}D_3xJ&d@=6j(6BZKPl?!k1?!`f3z&a zR4ZF60Mx7oBxLSxGuzA*Dy5n-d2K=+)6VMZh_0KetK|{e;E{8NJJ!)=_E~1uu=A=r zrn&gh)h*SFhsQJo!f+wKMIE;-EOaMSMB@aXRU(UcnJhZW^B^mgs|M9@5WF@s6B0p& zm#CTz)yiQCgURE{%hjxHcJ6G&>G9i`7MyftL!QQd5 z@RflRs?7)99?X`kHNt>W3l7YqscBpi*R2+fsgABor>KVOu(i(`03aytf2UA!&SC9v z!E}whj#^9~=XHMinFZ;6UOJjo=mmNaWkv~nC=qH9$s-8roGeyaW-E~SzZ3Gg>j zZ8}<320rg4=$`M0nxN!w(PtHUjeeU?MvYgWKZ6kkzABK;vMN0|U;X9abJleJA(xy<}5h5P(5 z{RzAFPvMnX2m0yH0Jn2Uo-p`daE|(O`YQiC#jB8;6bVIUf?SY(k$#C0`d6qT`>Xe0+0}Oj0=F&*D;PVe=Z<=0AGI<6$gYLwa#r` zm449x*fU;_+J>Mz!wa;T-wldoBB%&OEMJgtm#oaI60TSYCy7;+$5?q!zi5K`u66Wq zvg)Fx$s`V3Em{=OEY{3lmh_7|08ykS&U9w!kp@Ctuzqe1JFOGz6%i5}Kmm9>^=gih z?kRxqLA<3@e=}G4R_?phW{4DVr?`tPfyZSN@R=^;P;?!2bh~F1I|fB7P=V=9a6XU5 z<#0f>RS0O&rhc&nTRFOW7&QhevP0#>j0eq<1@D5yAlgMl5n&O9X|Vq}%RX}iNyRFF z7sX&u#6?E~bm~N|z&YikXC=I0E*8Z$v7PtWfjy)$e_Ez25fnR1Q=q1`;U!~U>|&YS zaOS8y!^ORmr2L4ik!IYR8@Dcx8MTC=(b4P6iE5CnrbI~7j7DmM8em$!da&D!6Xu)!vKPdLG z9f#)se|6=5yOCe)N6xDhPI!m81*dNe7u985zi%IVfOfJh69+#ag4ELzGne?o`eA`42K4T)h3S+s)5IT97%O>du- z0U54L8m4}rkRQ?QBfJ%DLssy^+a7Ajw;0&`NOTY4o;0-ivm9 zBz1C%nr_hQ)X)^QM6T1?=yeLkuG9Lf50(eH}`tFye;01&(p?8i+6h};VV-2B~qdxeC#=X z(JLlzy&fHkyi9Ksbcs~&r^%lh^2COldLz^H@X!s~mr9Dr6z!j+4?zkD@Ls7F8(t(f z9`U?P$Lmn*Y{K}aR4N&1N=?xtQ1%jqf1~pJyQ4SgBrEtR`j4lQuh7cqP49Em5cO=I zB(He2`iPN5M=Y0}h(IU$37ANTGx&|b-u1BYA*#dE(L-lptoOpo&th~E)_)y-`6kSH z3vvyVrcBwW^_XYReJ=JYd9OBQrzv;f2AQdZH#$Y{Y+Oa33M70XFI((fs;mB4e`<<{ ze4dv2B0V_?Ytsi>>g%qs*}oDGd5d(RNZ*6?7qNbdp7wP4T72=F&r?Ud#kZr8Ze5tB z_oNb7{G+(o2ajL$!69FW@jjPQ2a5C)m!MKKRirC$_VYIuVQCpf9rIms0GRDf)8AH${I`q^~5rjot@#3$2#zT2f`(N^P7Z;6(@EK$q*Jgif00I6*^ZGV+XB5uw*1R-@23yTw&WKD{s1;HTL;dO)%5i#`dc6b7;5@^{KU%N|A-$zsYw4)7LA{3`Zp>1 z-?K9_IE&z)dayUM)wd8K^29m-l$lFhi$zj0l!u~4;VGR6Y!?MAfBC^?QD53hy6VdD z@eUZIui}~L%#SmajaRq1J|#> z4m=o$vZ*34=ZWK2!QMNEcp2Lbc5N1q!lEDq(bz0b;WI9;e>l=CG9^n#ro`w>_0F$Q zfZ={2QyTkfByC&gy;x!r*NyXXbk=a%~~(#K?< zTke0HuF5{Q+~?@!KDXR|g+43$+;ab`^flS%miup_0OUTm=nIc%d5nLP)i308PIjl_YMF6cpQ__6&$n6it8K- z8PIjl_YMF6cpQ_!r)L8IivW`WdK8mBs6PXdjR2DYdK8nCs73=4j{uVadK8oNjwX|E wpAeHLsTu^*Y>Trk?aBtSQ(D-o$(D8Px^?ZI-PUB? z*1fv!{YdHme3Fc8%cR@*@zc5A_nq&2=R47Hp@$-JF4Fz*;SLw5}K^y>s-s;V!}b2i=5=M- zComP?ju>8Fe@=H@rlwe1l`J*6BTTo`9b$zjQ@HxrAhp0D#u?M~TxGC_!?ccCHCjt| zF*PgJf@kJB`|Ml}cmsyrAjO#Kjr^E5p29w+#>$C`Q|54BoDv$fQ9D?3n32P9LPMIzu?LjNqggOH=1@T{9bMn*u8(GI z!;MLTtFPHal^S>VcJdiYqX0VU|Rn@A}C1xOlxCribxes0~+n2 z6qDaIA2$?e`opx3_KW!rAgbpzU)gFdjAKXh|5w``#F0R|c)Y)Du0_Ihhz^S?k^pk% zP>9|pIDx)xHH^_~+aA=^$M!<8K~Hy(71nJGf6`HnjtS=4X4=Hk^O71oNia2V{HUCC zoN3RSBS?mZCLw;l4W4a+D8qc)XJS`pUJ5X-f^1ytxwr`@si$lAE?{4G|o; zO0l>`rr?;~c;{ZEFJ!!3=7=FdGJ?Q^xfNQh4A?i;IJ4}B+A?4olTK(fN++3CRBP97 ze~lG9h%oegkn)lpW-4F8o2`*WW0mZHwHez`ko@>U1_;EC_6ig|Drn@=DMV9YEUSCa zIf$kHei3(u#zm9I!Jf(4t`Vm1lltJ&lVHy(eIXE8sy9sUpmz%I_gA#8x^Zv8%w?r2 z{GdkX1SkzRIr>prRK@rqn9j2wG|rUvf6PJbbin=yy-TAXrguvzN8jL$hUrIXzr^s5 zVM?H4;eM-QeRFr06@ifV(ocvk?_)~N@1c2ien56UjWXid6W%6ievIh)>dk|rIs##^kY67ib8Kw%#-oVFaXG7$ERyA9(NSJUvWiOA5H(!{uOpcW zg&-?iqPhds%3%tFspHDqqr;A!e@B#iPQjHd=c>N1LoOEGRehVoPOdxJ>b6>yc#o#+ zl8s8!(|NMeqjsy@0x{8^j0d00SqRZjp{Kj)&4UHYGxG+z9b-)72I*&J70?+8e?p_@ z=>-(>l6z5vYlP~<2%DU02b!mA{7mS)NS_eLe=t)sm&+Pmk?asOEKlkPQ)EUvvfC=;4M&*|I!w}(@V_)eUKLA_t^%`o z0PM9LV|UKTLnk|?M3u!|f2S0?UqZsEIH9*NJS-8lzu;A6-rr-ot=dg9SASoluZUkFH$7X; zP=?kYX!K?JL-b~<#7wU;b;eS)O;@?h%sPPk{4xEBxb{!sm0AY|f9cNvx6>$3F!*0c z75H=dy8JvTyO8}g1w{$9T$p~5en}AeSLoCF>_RT9YPMpChUjl310o*$QocjbH& zbnwg#gssR#jDVN{uEi3n(PZ%PFZ|6J2 z5_rBf0-u>e4sFe0*Km49ATi7>Kn0f9!uc|rRMR1Dtt6m1LW8^>qFlo}h$@br=Rmpi z;mI&>OF64Be{dVeHI8utrh)v^wsZ0jii%x8UgZ8TC%K~@I(4E};GFW&(;WVov}3%H zH;IhRkfD^(vt^DjZz(MyHLZxv8}qzPc(%itBkBwf_fC~sDBgh<3XAv5cxxfF3<2U! z03Xe&z`is!JDHbe;mNmfkH+_LFE*I2^mdL@7(@9DfAcP6O04V-ko;Rpgp<%Cj5r8Z zd0`sXoIjV$j)--;jA6Zy^D5&5v$o^>e%>Q?9GLm{i~p^lAn!%ZtF$I~>39XVZxk0b zROh^Bk9cE0AJBLozZIEmy7xG(yHWGztvfnr0(2ro1%>zsGMS^EMu+S$r=_;9 zWwZkgf7Q7`H9sLf2Go^Xy6&h~a&%s2_T@_Csf19MntF$aVFiFkvE3_hUg(B@&Xw@YJ zpL$wNYf78=0c@!QU6_a$>CPiXT7QAGDM}7Z(0z#_ZA=fmLUj{2z7@Ypo71UDy8GHr z-&TLKf6a5WCf@Adle3VglBt4>Z>;xF}}-S~B7<(%B;Y z0QR55{z-buw>8ilNM3u6I+D$S%?)(p>=eBx-HpvZj{7c*_?K=d()*7q?93us}1dq%FAFYLsW8ZTQ_XZLh`P2*6(NgS}qGcfGXVWpwsp#Rs}IuKbk*`2}&) zI^Vsk6S&Q4@oYS?dJ`NwMVBs6f57+RxdqVub#PvMu?$=^OJy5xEl0<5SLsSRy%%a0 zi}Y#1-F3m;Ieh#Y12UgW?-R)|eX>ZuF-2cc!1>~NS|XSF-6In>zBoZg+ml!6%fk7U zw0LHcz8VQk(jOJ+Yu)|^|15ufl$KQd_1eUZZzj`aC%umU6F1&D5XVWce_wAe(qCSZ zpX-QF4e{EmEVN9~6%bR5U*UT{eMHfcUo`jw*u?4r2s_$`}U{?NjvEm(u&<>B|%mq$Q3weshxk z76<``8vh{+nX`@9CB6IE&z)I%IFjR^LH{s1p|eppv=x za(g_jLU|xjWMAn-V7th$f({|LG8zzIE0g?cyW;%Dmtv%C+0@xVxPE^ zyZzi9P%JAD6ynwHptuzP`Kox7*9h7XSMonCalv;Md0i9Vb-c*!f0ubfk?&T&T}AHh z4m8Bz{JllKcdNg?D^%a5MFQ;#1z|*}H^qHLzW)L}wp?2tY7RejtSh8<;Zw)QGJYUm z|MbTxyj*McKlStlT9I5XlSWtQGN&-LTr2XyNU+`490rg?LYLMRnz-@oKqT1hpCGqP zyRXt4=_Woj$%n5ee<3zhLF>5>`?m9a#xQH+Jk_+|RM8Vi;2*XbK- zEL6sCpaGPzP>k8f4Kh|##_imt#zJMB;ir|JrMPGW`rityK1vHXMLy18%qmMQAm4WZ zP)i30KR&5vs15)C+8dM66&$k~i|ZT;KR&5vs15)C+8dJ(sAmGPijyIz6_bsqKLSFH zlOd=TljEpH0>h4zA*dCTK&emy#FCRCs1=i^sZ9bFmXjf<6_X39E(XY)00000#N437 diff --git a/gradle/wrapper/gradle-wrapper.properties b/gradle/wrapper/gradle-wrapper.properties index c7f182843385d..39a291b258efb 100644 --- a/gradle/wrapper/gradle-wrapper.properties +++ b/gradle/wrapper/gradle-wrapper.properties @@ -11,7 +11,7 @@ distributionBase=GRADLE_USER_HOME distributionPath=wrapper/dists -distributionUrl=https\://services.gradle.org/distributions/gradle-8.9-all.zip +distributionUrl=https\://services.gradle.org/distributions/gradle-8.10-all.zip zipStoreBase=GRADLE_USER_HOME zipStorePath=wrapper/dists -distributionSha256Sum=258e722ec21e955201e31447b0aed14201765a3bfbae296a46cf60b70e66db70 +distributionSha256Sum=682b4df7fe5accdca84a4d1ef6a3a6ab096b3efd5edf7de2bd8c758d95a93703 From 1717b551db0d6e5d22aff89bf6fd9873d8d4ec20 Mon Sep 17 00:00:00 2001 From: Sachin Kale Date: Thu, 15 Aug 2024 15:21:17 +0530 Subject: [PATCH 35/40] Add timestamp pinning service and scheduler to update in-memory state (#15180) --------- Signed-off-by: Sachin Kale Co-authored-by: Sachin Kale --- .../RemoteStorePinnedTimestampsIT.java | 86 +++++ .../common/settings/ClusterSettings.java | 3 + .../remote/model/RemotePinnedTimestamps.java | 144 ++++++++ .../RemoteStorePinnedTimestampsBlobStore.java | 43 +++ .../main/java/org/opensearch/node/Node.java | 23 +- .../RemoteStorePinnedTimestampService.java | 321 ++++++++++++++++++ .../snapshots/SnapshotsService.java | 6 +- .../model/RemotePinnedTimestampsTests.java | 101 ++++++ .../snapshots/SnapshotResiliencyTests.java | 3 +- 9 files changed, 727 insertions(+), 3 deletions(-) create mode 100644 server/src/internalClusterTest/java/org/opensearch/remotestore/RemoteStorePinnedTimestampsIT.java create mode 100644 server/src/main/java/org/opensearch/gateway/remote/model/RemotePinnedTimestamps.java create mode 100644 server/src/main/java/org/opensearch/gateway/remote/model/RemoteStorePinnedTimestampsBlobStore.java create mode 100644 server/src/main/java/org/opensearch/node/remotestore/RemoteStorePinnedTimestampService.java create mode 100644 server/src/test/java/org/opensearch/gateway/remote/model/RemotePinnedTimestampsTests.java diff --git a/server/src/internalClusterTest/java/org/opensearch/remotestore/RemoteStorePinnedTimestampsIT.java b/server/src/internalClusterTest/java/org/opensearch/remotestore/RemoteStorePinnedTimestampsIT.java new file mode 100644 index 0000000000000..0bb53309f7a78 --- /dev/null +++ b/server/src/internalClusterTest/java/org/opensearch/remotestore/RemoteStorePinnedTimestampsIT.java @@ -0,0 +1,86 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.remotestore; + +import org.opensearch.common.collect.Tuple; +import org.opensearch.common.unit.TimeValue; +import org.opensearch.core.action.ActionListener; +import org.opensearch.node.remotestore.RemoteStorePinnedTimestampService; +import org.opensearch.test.OpenSearchIntegTestCase; + +import java.util.Set; + +@OpenSearchIntegTestCase.ClusterScope(scope = OpenSearchIntegTestCase.Scope.TEST, numDataNodes = 0) +public class RemoteStorePinnedTimestampsIT extends RemoteStoreBaseIntegTestCase { + static final String INDEX_NAME = "remote-store-test-idx-1"; + + ActionListener noOpActionListener = new ActionListener<>() { + @Override + public void onResponse(Void unused) {} + + @Override + public void onFailure(Exception e) {} + }; + + public void testTimestampPinUnpin() throws Exception { + prepareCluster(1, 1, INDEX_NAME, 0, 2); + ensureGreen(INDEX_NAME); + + RemoteStorePinnedTimestampService remoteStorePinnedTimestampService = internalCluster().getInstance( + RemoteStorePinnedTimestampService.class, + primaryNodeName(INDEX_NAME) + ); + + Tuple> pinnedTimestampWithFetchTimestamp = RemoteStorePinnedTimestampService.getPinnedTimestamps(); + long lastFetchTimestamp = pinnedTimestampWithFetchTimestamp.v1(); + assertEquals(-1L, lastFetchTimestamp); + assertEquals(Set.of(), pinnedTimestampWithFetchTimestamp.v2()); + + assertThrows( + IllegalArgumentException.class, + () -> remoteStorePinnedTimestampService.pinTimestamp(1234L, "ss1", noOpActionListener) + ); + + long timestamp1 = System.currentTimeMillis() + 30000L; + long timestamp2 = System.currentTimeMillis() + 60000L; + long timestamp3 = System.currentTimeMillis() + 900000L; + remoteStorePinnedTimestampService.pinTimestamp(timestamp1, "ss2", noOpActionListener); + remoteStorePinnedTimestampService.pinTimestamp(timestamp2, "ss3", noOpActionListener); + remoteStorePinnedTimestampService.pinTimestamp(timestamp3, "ss4", noOpActionListener); + + remoteStorePinnedTimestampService.setPinnedTimestampsSchedulerInterval(TimeValue.timeValueSeconds(1)); + + assertBusy(() -> { + Tuple> pinnedTimestampWithFetchTimestamp_2 = RemoteStorePinnedTimestampService.getPinnedTimestamps(); + long lastFetchTimestamp_2 = pinnedTimestampWithFetchTimestamp_2.v1(); + assertTrue(lastFetchTimestamp_2 != -1); + assertEquals(Set.of(timestamp1, timestamp2, timestamp3), pinnedTimestampWithFetchTimestamp_2.v2()); + }); + + remoteStorePinnedTimestampService.setPinnedTimestampsSchedulerInterval(TimeValue.timeValueMinutes(3)); + + // This should be a no-op as pinning entity is different + remoteStorePinnedTimestampService.unpinTimestamp(timestamp1, "no-snapshot", noOpActionListener); + // Unpinning already pinned entity + remoteStorePinnedTimestampService.unpinTimestamp(timestamp2, "ss3", noOpActionListener); + // Adding different entity to already pinned timestamp + remoteStorePinnedTimestampService.pinTimestamp(timestamp3, "ss5", noOpActionListener); + + remoteStorePinnedTimestampService.setPinnedTimestampsSchedulerInterval(TimeValue.timeValueSeconds(1)); + + assertBusy(() -> { + Tuple> pinnedTimestampWithFetchTimestamp_3 = RemoteStorePinnedTimestampService.getPinnedTimestamps(); + long lastFetchTimestamp_3 = pinnedTimestampWithFetchTimestamp_3.v1(); + assertTrue(lastFetchTimestamp_3 != -1); + assertEquals(Set.of(timestamp1, timestamp3), pinnedTimestampWithFetchTimestamp_3.v2()); + }); + + remoteStorePinnedTimestampService.setPinnedTimestampsSchedulerInterval(TimeValue.timeValueMinutes(3)); + } +} diff --git a/server/src/main/java/org/opensearch/common/settings/ClusterSettings.java b/server/src/main/java/org/opensearch/common/settings/ClusterSettings.java index a73e5d44b7e02..7baae17dd77cd 100644 --- a/server/src/main/java/org/opensearch/common/settings/ClusterSettings.java +++ b/server/src/main/java/org/opensearch/common/settings/ClusterSettings.java @@ -142,6 +142,7 @@ import org.opensearch.node.Node.DiscoverySettings; import org.opensearch.node.NodeRoleSettings; import org.opensearch.node.remotestore.RemoteStoreNodeService; +import org.opensearch.node.remotestore.RemoteStorePinnedTimestampService; import org.opensearch.node.resource.tracker.ResourceTrackerSettings; import org.opensearch.persistent.PersistentTasksClusterService; import org.opensearch.persistent.decider.EnableAssignmentDecider; @@ -760,6 +761,8 @@ public void apply(Settings value, Settings current, Settings previous) { RemoteStoreSettings.CLUSTER_REMOTE_STORE_TRANSLOG_METADATA, SearchService.CLUSTER_ALLOW_DERIVED_FIELD_SETTING, + RemoteStorePinnedTimestampService.CLUSTER_REMOTE_STORE_PINNED_TIMESTAMP_SCHEDULER_INTERVAL, + // Composite index settings CompositeIndexSettings.STAR_TREE_INDEX_ENABLED_SETTING, diff --git a/server/src/main/java/org/opensearch/gateway/remote/model/RemotePinnedTimestamps.java b/server/src/main/java/org/opensearch/gateway/remote/model/RemotePinnedTimestamps.java new file mode 100644 index 0000000000000..030491cf8b7b9 --- /dev/null +++ b/server/src/main/java/org/opensearch/gateway/remote/model/RemotePinnedTimestamps.java @@ -0,0 +1,144 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.gateway.remote.model; + +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.opensearch.common.io.Streams; +import org.opensearch.common.remote.BlobPathParameters; +import org.opensearch.common.remote.RemoteWriteableBlobEntity; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.core.common.io.stream.StreamOutput; +import org.opensearch.core.common.io.stream.Writeable; +import org.opensearch.core.compress.Compressor; +import org.opensearch.index.remote.RemoteStoreUtils; +import org.opensearch.repositories.blobstore.ChecksumWritableBlobStoreFormat; + +import java.io.IOException; +import java.io.InputStream; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; + +import static org.opensearch.gateway.remote.RemoteClusterStateUtils.DELIMITER; + +/** + * Wrapper class for uploading/downloading {@link RemotePinnedTimestamps} to/from remote blob store + * + * @opensearch.internal + */ +public class RemotePinnedTimestamps extends RemoteWriteableBlobEntity { + private static final Logger logger = LogManager.getLogger(RemotePinnedTimestamps.class); + + /** + * Represents a collection of pinned timestamps and their associated pinning entities. + * This class is thread-safe and implements the Writeable interface for serialization. + */ + public static class PinnedTimestamps implements Writeable { + private final Map> pinnedTimestampPinningEntityMap; + + public PinnedTimestamps(Map> pinnedTimestampPinningEntityMap) { + this.pinnedTimestampPinningEntityMap = new ConcurrentHashMap<>(pinnedTimestampPinningEntityMap); + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + out.writeMap(pinnedTimestampPinningEntityMap, StreamOutput::writeLong, StreamOutput::writeStringCollection); + } + + public static PinnedTimestamps readFrom(StreamInput in) throws IOException { + return new PinnedTimestamps(in.readMap(StreamInput::readLong, StreamInput::readStringList)); + } + + /** + * Pins a timestamp against a pinning entity. + * + * @param timestamp The timestamp to pin. + * @param pinningEntity The entity pinning the timestamp. + */ + public void pin(Long timestamp, String pinningEntity) { + logger.debug("Pinning timestamp = {} against entity = {}", timestamp, pinningEntity); + pinnedTimestampPinningEntityMap.computeIfAbsent(timestamp, k -> new ArrayList<>()).add(pinningEntity); + } + + /** + * Unpins a timestamp for a specific pinning entity. + * + * @param timestamp The timestamp to unpin. + * @param pinningEntity The entity unpinning the timestamp. + */ + public void unpin(Long timestamp, String pinningEntity) { + logger.debug("Unpinning timestamp = {} against entity = {}", timestamp, pinningEntity); + if (pinnedTimestampPinningEntityMap.containsKey(timestamp) == false + || pinnedTimestampPinningEntityMap.get(timestamp).contains(pinningEntity) == false) { + logger.warn("Timestamp: {} is not pinned by entity: {}", timestamp, pinningEntity); + } + pinnedTimestampPinningEntityMap.compute(timestamp, (k, v) -> { + v.remove(pinningEntity); + return v.isEmpty() ? null : v; + }); + } + + public Map> getPinnedTimestampPinningEntityMap() { + return new HashMap<>(pinnedTimestampPinningEntityMap); + } + } + + public static final String PINNED_TIMESTAMPS = "pinned_timestamps"; + public static final ChecksumWritableBlobStoreFormat PINNED_TIMESTAMPS_FORMAT = new ChecksumWritableBlobStoreFormat<>( + PINNED_TIMESTAMPS, + PinnedTimestamps::readFrom + ); + + private PinnedTimestamps pinnedTimestamps; + + public RemotePinnedTimestamps(String clusterUUID, Compressor compressor) { + super(clusterUUID, compressor); + pinnedTimestamps = new PinnedTimestamps(new HashMap<>()); + } + + @Override + public BlobPathParameters getBlobPathParameters() { + return new BlobPathParameters(List.of(PINNED_TIMESTAMPS), PINNED_TIMESTAMPS); + } + + @Override + public String getType() { + return PINNED_TIMESTAMPS; + } + + @Override + public String generateBlobFileName() { + return this.blobFileName = String.join(DELIMITER, PINNED_TIMESTAMPS, RemoteStoreUtils.invertLong(System.currentTimeMillis())); + } + + @Override + public InputStream serialize() throws IOException { + return PINNED_TIMESTAMPS_FORMAT.serialize(pinnedTimestamps, generateBlobFileName(), getCompressor()).streamInput(); + } + + @Override + public PinnedTimestamps deserialize(InputStream inputStream) throws IOException { + return PINNED_TIMESTAMPS_FORMAT.deserialize(blobName, Streams.readFully(inputStream)); + } + + public void setBlobFileName(String blobFileName) { + this.blobFileName = blobFileName; + } + + public void setPinnedTimestamps(PinnedTimestamps pinnedTimestamps) { + this.pinnedTimestamps = pinnedTimestamps; + } + + public PinnedTimestamps getPinnedTimestamps() { + return pinnedTimestamps; + } +} diff --git a/server/src/main/java/org/opensearch/gateway/remote/model/RemoteStorePinnedTimestampsBlobStore.java b/server/src/main/java/org/opensearch/gateway/remote/model/RemoteStorePinnedTimestampsBlobStore.java new file mode 100644 index 0000000000000..2a65dd993d0af --- /dev/null +++ b/server/src/main/java/org/opensearch/gateway/remote/model/RemoteStorePinnedTimestampsBlobStore.java @@ -0,0 +1,43 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.gateway.remote.model; + +import org.opensearch.common.blobstore.BlobPath; +import org.opensearch.common.remote.RemoteWriteableBlobEntity; +import org.opensearch.common.remote.RemoteWriteableEntityBlobStore; +import org.opensearch.index.translog.transfer.BlobStoreTransferService; +import org.opensearch.repositories.blobstore.BlobStoreRepository; +import org.opensearch.threadpool.ThreadPool; + +/** + * Extends the RemoteClusterStateBlobStore to support {@link RemotePinnedTimestamps} + */ +public class RemoteStorePinnedTimestampsBlobStore extends RemoteWriteableEntityBlobStore< + RemotePinnedTimestamps.PinnedTimestamps, + RemotePinnedTimestamps> { + + public static final String PINNED_TIMESTAMPS_PATH_TOKEN = "pinned_timestamps"; + private final BlobStoreRepository blobStoreRepository; + + public RemoteStorePinnedTimestampsBlobStore( + BlobStoreTransferService blobStoreTransferService, + BlobStoreRepository blobStoreRepository, + String clusterName, + ThreadPool threadPool, + String executor + ) { + super(blobStoreTransferService, blobStoreRepository, clusterName, threadPool, executor, PINNED_TIMESTAMPS_PATH_TOKEN); + this.blobStoreRepository = blobStoreRepository; + } + + @Override + public BlobPath getBlobPathForUpload(final RemoteWriteableBlobEntity obj) { + return blobStoreRepository.basePath().add(PINNED_TIMESTAMPS_PATH_TOKEN); + } +} diff --git a/server/src/main/java/org/opensearch/node/Node.java b/server/src/main/java/org/opensearch/node/Node.java index 409f84354a8b1..1a9b233b387b2 100644 --- a/server/src/main/java/org/opensearch/node/Node.java +++ b/server/src/main/java/org/opensearch/node/Node.java @@ -185,6 +185,7 @@ import org.opensearch.monitor.fs.FsProbe; import org.opensearch.monitor.jvm.JvmInfo; import org.opensearch.node.remotestore.RemoteStoreNodeService; +import org.opensearch.node.remotestore.RemoteStorePinnedTimestampService; import org.opensearch.node.resource.tracker.NodeResourceUsageTracker; import org.opensearch.persistent.PersistentTasksClusterService; import org.opensearch.persistent.PersistentTasksExecutor; @@ -810,6 +811,18 @@ protected Node( remoteIndexPathUploader = null; remoteClusterStateCleanupManager = null; } + final RemoteStorePinnedTimestampService remoteStorePinnedTimestampService; + if (isRemoteStoreAttributePresent(settings)) { + remoteStorePinnedTimestampService = new RemoteStorePinnedTimestampService( + repositoriesServiceReference::get, + settings, + threadPool, + clusterService + ); + resourcesToClose.add(remoteStorePinnedTimestampService); + } else { + remoteStorePinnedTimestampService = null; + } // collect engine factory providers from plugins final Collection enginePlugins = pluginsService.filterPlugins(EnginePlugin.class); @@ -1173,7 +1186,8 @@ protected Node( clusterModule.getIndexNameExpressionResolver(), repositoryService, transportService, - actionModule.getActionFilters() + actionModule.getActionFilters(), + remoteStorePinnedTimestampService ); SnapshotShardsService snapshotShardsService = new SnapshotShardsService( settings, @@ -1426,6 +1440,7 @@ protected Node( b.bind(MetricsRegistry.class).toInstance(metricsRegistry); b.bind(RemoteClusterStateService.class).toProvider(() -> remoteClusterStateService); b.bind(RemoteIndexPathUploader.class).toProvider(() -> remoteIndexPathUploader); + b.bind(RemoteStorePinnedTimestampService.class).toProvider(() -> remoteStorePinnedTimestampService); b.bind(RemoteClusterStateCleanupManager.class).toProvider(() -> remoteClusterStateCleanupManager); b.bind(PersistedStateRegistry.class).toInstance(persistedStateRegistry); b.bind(SegmentReplicationStatsTracker.class).toInstance(segmentReplicationStatsTracker); @@ -1581,6 +1596,12 @@ public Node start() throws NodeValidationException { if (remoteIndexPathUploader != null) { remoteIndexPathUploader.start(); } + final RemoteStorePinnedTimestampService remoteStorePinnedTimestampService = injector.getInstance( + RemoteStorePinnedTimestampService.class + ); + if (remoteStorePinnedTimestampService != null) { + remoteStorePinnedTimestampService.start(); + } // Load (and maybe upgrade) the metadata stored on disk final GatewayMetaState gatewayMetaState = injector.getInstance(GatewayMetaState.class); gatewayMetaState.start( diff --git a/server/src/main/java/org/opensearch/node/remotestore/RemoteStorePinnedTimestampService.java b/server/src/main/java/org/opensearch/node/remotestore/RemoteStorePinnedTimestampService.java new file mode 100644 index 0000000000000..35730a75a8142 --- /dev/null +++ b/server/src/main/java/org/opensearch/node/remotestore/RemoteStorePinnedTimestampService.java @@ -0,0 +1,321 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.node.remotestore; + +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.opensearch.cluster.ClusterName; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.blobstore.BlobMetadata; +import org.opensearch.common.blobstore.BlobPath; +import org.opensearch.common.collect.Tuple; +import org.opensearch.common.settings.Setting; +import org.opensearch.common.settings.Settings; +import org.opensearch.common.unit.TimeValue; +import org.opensearch.common.util.concurrent.AbstractAsyncTask; +import org.opensearch.core.action.ActionListener; +import org.opensearch.gateway.remote.model.RemotePinnedTimestamps; +import org.opensearch.gateway.remote.model.RemotePinnedTimestamps.PinnedTimestamps; +import org.opensearch.gateway.remote.model.RemoteStorePinnedTimestampsBlobStore; +import org.opensearch.index.translog.transfer.BlobStoreTransferService; +import org.opensearch.node.Node; +import org.opensearch.repositories.RepositoriesService; +import org.opensearch.repositories.Repository; +import org.opensearch.repositories.blobstore.BlobStoreRepository; +import org.opensearch.threadpool.ThreadPool; + +import java.io.Closeable; +import java.io.IOException; +import java.util.HashMap; +import java.util.List; +import java.util.Set; +import java.util.concurrent.Semaphore; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import java.util.function.Consumer; +import java.util.function.Supplier; +import java.util.stream.Collectors; + +/** + * Service for managing pinned timestamps in a remote store. + * This service handles pinning and unpinning of timestamps, as well as periodic updates of the pinned timestamps set. + * + * @opensearch.internal + */ +public class RemoteStorePinnedTimestampService implements Closeable { + private static final Logger logger = LogManager.getLogger(RemoteStorePinnedTimestampService.class); + private static Tuple> pinnedTimestampsSet = new Tuple<>(-1L, Set.of()); + public static final int PINNED_TIMESTAMP_FILES_TO_KEEP = 5; + + private final Supplier repositoriesService; + private final Settings settings; + private final ThreadPool threadPool; + private final ClusterService clusterService; + private BlobStoreRepository blobStoreRepository; + private BlobStoreTransferService blobStoreTransferService; + private RemoteStorePinnedTimestampsBlobStore pinnedTimestampsBlobStore; + private AsyncUpdatePinnedTimestampTask asyncUpdatePinnedTimestampTask; + private volatile TimeValue pinnedTimestampsSchedulerInterval; + private final Semaphore updateTimetampPinningSemaphore = new Semaphore(1); + + /** + * Controls pinned timestamp scheduler interval + */ + public static final Setting CLUSTER_REMOTE_STORE_PINNED_TIMESTAMP_SCHEDULER_INTERVAL = Setting.timeSetting( + "cluster.remote_store.pinned_timestamps.scheduler_interval", + TimeValue.timeValueMinutes(3), + TimeValue.timeValueMinutes(1), + Setting.Property.NodeScope + ); + + /** + * Controls allowed timestamp values to be pinned from past + */ + public static final Setting CLUSTER_REMOTE_STORE_PINNED_TIMESTAMP_LOOKBACK_INTERVAL = Setting.timeSetting( + "cluster.remote_store.pinned_timestamps.lookback_interval", + TimeValue.timeValueMinutes(1), + TimeValue.timeValueMinutes(1), + TimeValue.timeValueMinutes(5), + Setting.Property.NodeScope + ); + + public RemoteStorePinnedTimestampService( + Supplier repositoriesService, + Settings settings, + ThreadPool threadPool, + ClusterService clusterService + ) { + this.repositoriesService = repositoriesService; + this.settings = settings; + this.threadPool = threadPool; + this.clusterService = clusterService; + + pinnedTimestampsSchedulerInterval = CLUSTER_REMOTE_STORE_PINNED_TIMESTAMP_SCHEDULER_INTERVAL.get(settings); + } + + /** + * Starts the RemoteStorePinnedTimestampService. + * This method validates the remote store configuration, initializes components, + * and starts the asynchronous update task. + */ + public void start() { + validateRemoteStoreConfiguration(); + initializeComponents(); + startAsyncUpdateTask(); + } + + private void validateRemoteStoreConfiguration() { + final String remoteStoreRepo = settings.get( + Node.NODE_ATTRIBUTES.getKey() + RemoteStoreNodeAttribute.REMOTE_STORE_SEGMENT_REPOSITORY_NAME_ATTRIBUTE_KEY + ); + assert remoteStoreRepo != null : "Remote Segment Store repository is not configured"; + final Repository repository = repositoriesService.get().repository(remoteStoreRepo); + assert repository instanceof BlobStoreRepository : "Repository should be instance of BlobStoreRepository"; + blobStoreRepository = (BlobStoreRepository) repository; + } + + private void initializeComponents() { + String clusterName = ClusterName.CLUSTER_NAME_SETTING.get(settings).value(); + blobStoreTransferService = new BlobStoreTransferService(blobStoreRepository.blobStore(), this.threadPool); + pinnedTimestampsBlobStore = new RemoteStorePinnedTimestampsBlobStore( + blobStoreTransferService, + blobStoreRepository, + clusterName, + this.threadPool, + ThreadPool.Names.REMOTE_STATE_READ + ); + } + + private void startAsyncUpdateTask() { + asyncUpdatePinnedTimestampTask = new AsyncUpdatePinnedTimestampTask(logger, threadPool, pinnedTimestampsSchedulerInterval, true); + } + + /** + * Pins a timestamp in the remote store. + * + * @param timestamp The timestamp to be pinned + * @param pinningEntity The entity responsible for pinning the timestamp + * @param listener A listener to be notified when the pinning operation completes + * @throws IllegalArgumentException If the timestamp is less than the current time minus one second + */ + public void pinTimestamp(long timestamp, String pinningEntity, ActionListener listener) { + // If a caller uses current system time to pin the timestamp, following check will almost always fail. + // So, we allow pinning timestamp in the past upto some buffer + long lookbackIntervalInMills = CLUSTER_REMOTE_STORE_PINNED_TIMESTAMP_LOOKBACK_INTERVAL.get(settings).millis(); + if (timestamp < TimeUnit.NANOSECONDS.toMillis(System.nanoTime()) - lookbackIntervalInMills) { + throw new IllegalArgumentException( + "Timestamp to be pinned is less than current timestamp - value of cluster.remote_store.pinned_timestamps.lookback_interval" + ); + } + updatePinning(pinnedTimestamps -> pinnedTimestamps.pin(timestamp, pinningEntity), listener); + } + + /** + * Unpins a timestamp from the remote store. + * + * @param timestamp The timestamp to be unpinned + * @param pinningEntity The entity responsible for unpinning the timestamp + * @param listener A listener to be notified when the unpinning operation completes + */ + public void unpinTimestamp(long timestamp, String pinningEntity, ActionListener listener) { + updatePinning(pinnedTimestamps -> pinnedTimestamps.unpin(timestamp, pinningEntity), listener); + } + + private void updatePinning(Consumer updateConsumer, ActionListener listener) { + RemotePinnedTimestamps remotePinnedTimestamps = new RemotePinnedTimestamps( + clusterService.state().metadata().clusterUUID(), + blobStoreRepository.getCompressor() + ); + BlobPath path = pinnedTimestampsBlobStore.getBlobPathForUpload(remotePinnedTimestamps); + try { + if (updateTimetampPinningSemaphore.tryAcquire(10, TimeUnit.MINUTES)) { + ActionListener semaphoreAwareListener = ActionListener.runBefore(listener, updateTimetampPinningSemaphore::release); + ActionListener> listCallResponseListener = getListenerForListCallResponse( + remotePinnedTimestamps, + updateConsumer, + semaphoreAwareListener + ); + blobStoreTransferService.listAllInSortedOrder( + path, + remotePinnedTimestamps.getType(), + Integer.MAX_VALUE, + listCallResponseListener + ); + } else { + throw new TimeoutException("Timed out while waiting to acquire lock in updatePinning"); + } + } catch (InterruptedException | TimeoutException e) { + listener.onFailure(e); + } + } + + private ActionListener> getListenerForListCallResponse( + RemotePinnedTimestamps remotePinnedTimestamps, + Consumer updateConsumer, + ActionListener listener + ) { + return ActionListener.wrap(blobMetadata -> { + PinnedTimestamps pinnedTimestamps = new PinnedTimestamps(new HashMap<>()); + if (blobMetadata.isEmpty() == false) { + pinnedTimestamps = readExistingPinnedTimestamps(blobMetadata.get(0).name(), remotePinnedTimestamps); + } + updateConsumer.accept(pinnedTimestamps); + remotePinnedTimestamps.setPinnedTimestamps(pinnedTimestamps); + ActionListener writeCallResponseListener = getListenerForWriteCallResponse( + remotePinnedTimestamps, + blobMetadata, + listener + ); + pinnedTimestampsBlobStore.writeAsync(remotePinnedTimestamps, writeCallResponseListener); + }, listener::onFailure); + } + + private ActionListener getListenerForWriteCallResponse( + RemotePinnedTimestamps remotePinnedTimestamps, + List blobMetadata, + ActionListener listener + ) { + return ActionListener.wrap(unused -> { + // Delete older pinnedTimestamp files + if (blobMetadata.size() > PINNED_TIMESTAMP_FILES_TO_KEEP) { + List oldFilesToBeDeleted = blobMetadata.subList(PINNED_TIMESTAMP_FILES_TO_KEEP, blobMetadata.size()) + .stream() + .map(BlobMetadata::name) + .collect(Collectors.toList()); + try { + blobStoreTransferService.deleteBlobs( + pinnedTimestampsBlobStore.getBlobPathForUpload(remotePinnedTimestamps), + oldFilesToBeDeleted + ); + } catch (IOException e) { + logger.error("Exception while deleting stale pinned timestamps", e); + } + } + listener.onResponse(null); + }, listener::onFailure); + } + + private PinnedTimestamps readExistingPinnedTimestamps(String blobFilename, RemotePinnedTimestamps remotePinnedTimestamps) { + remotePinnedTimestamps.setBlobFileName(blobFilename); + remotePinnedTimestamps.setFullBlobName(pinnedTimestampsBlobStore.getBlobPathForUpload(remotePinnedTimestamps)); + try { + return pinnedTimestampsBlobStore.read(remotePinnedTimestamps); + } catch (IOException e) { + throw new RuntimeException("Failed to read existing pinned timestamps", e); + } + } + + @Override + public void close() throws IOException { + asyncUpdatePinnedTimestampTask.close(); + } + + // Visible for testing + public void setPinnedTimestampsSchedulerInterval(TimeValue pinnedTimestampsSchedulerInterval) { + this.pinnedTimestampsSchedulerInterval = pinnedTimestampsSchedulerInterval; + rescheduleAsyncUpdatePinnedTimestampTask(); + } + + private void rescheduleAsyncUpdatePinnedTimestampTask() { + if (pinnedTimestampsSchedulerInterval != null) { + pinnedTimestampsSet = new Tuple<>(-1L, Set.of()); + asyncUpdatePinnedTimestampTask.close(); + startAsyncUpdateTask(); + } + } + + public static Tuple> getPinnedTimestamps() { + return pinnedTimestampsSet; + } + + /** + * Inner class for asynchronously updating the pinned timestamp set. + */ + private final class AsyncUpdatePinnedTimestampTask extends AbstractAsyncTask { + private AsyncUpdatePinnedTimestampTask(Logger logger, ThreadPool threadPool, TimeValue interval, boolean autoReschedule) { + super(logger, threadPool, interval, autoReschedule); + rescheduleIfNecessary(); + } + + @Override + protected boolean mustReschedule() { + return true; + } + + @Override + protected void runInternal() { + long triggerTimestamp = System.currentTimeMillis(); + RemotePinnedTimestamps remotePinnedTimestamps = new RemotePinnedTimestamps( + clusterService.state().metadata().clusterUUID(), + blobStoreRepository.getCompressor() + ); + BlobPath path = pinnedTimestampsBlobStore.getBlobPathForUpload(remotePinnedTimestamps); + blobStoreTransferService.listAllInSortedOrder(path, remotePinnedTimestamps.getType(), 1, new ActionListener<>() { + @Override + public void onResponse(List blobMetadata) { + if (blobMetadata.isEmpty()) { + return; + } + PinnedTimestamps pinnedTimestamps = readExistingPinnedTimestamps(blobMetadata.get(0).name(), remotePinnedTimestamps); + logger.debug( + "Fetched pinned timestamps from remote store: {} - {}", + triggerTimestamp, + pinnedTimestamps.getPinnedTimestampPinningEntityMap().keySet() + ); + pinnedTimestampsSet = new Tuple<>(triggerTimestamp, pinnedTimestamps.getPinnedTimestampPinningEntityMap().keySet()); + } + + @Override + public void onFailure(Exception e) { + logger.error("Exception while listing pinned timestamp files", e); + } + }); + } + } +} diff --git a/server/src/main/java/org/opensearch/snapshots/SnapshotsService.java b/server/src/main/java/org/opensearch/snapshots/SnapshotsService.java index acc2dc83749cd..5e49208465dbb 100644 --- a/server/src/main/java/org/opensearch/snapshots/SnapshotsService.java +++ b/server/src/main/java/org/opensearch/snapshots/SnapshotsService.java @@ -92,6 +92,7 @@ import org.opensearch.core.index.Index; import org.opensearch.core.index.shard.ShardId; import org.opensearch.index.store.lockmanager.RemoteStoreLockManagerFactory; +import org.opensearch.node.remotestore.RemoteStorePinnedTimestampService; import org.opensearch.repositories.IndexId; import org.opensearch.repositories.RepositoriesService; import org.opensearch.repositories.Repository; @@ -180,6 +181,7 @@ public class SnapshotsService extends AbstractLifecycleComponent implements Clus private final UpdateSnapshotStatusAction updateSnapshotStatusHandler; private final TransportService transportService; + private final RemoteStorePinnedTimestampService remoteStorePinnedTimestampService; private final OngoingRepositoryOperations repositoryOperations = new OngoingRepositoryOperations(); @@ -208,7 +210,8 @@ public SnapshotsService( IndexNameExpressionResolver indexNameExpressionResolver, RepositoriesService repositoriesService, TransportService transportService, - ActionFilters actionFilters + ActionFilters actionFilters, + @Nullable RemoteStorePinnedTimestampService remoteStorePinnedTimestampService ) { this.clusterService = clusterService; this.indexNameExpressionResolver = indexNameExpressionResolver; @@ -216,6 +219,7 @@ public SnapshotsService( this.remoteStoreLockManagerFactory = new RemoteStoreLockManagerFactory(() -> repositoriesService); this.threadPool = transportService.getThreadPool(); this.transportService = transportService; + this.remoteStorePinnedTimestampService = remoteStorePinnedTimestampService; // The constructor of UpdateSnapshotStatusAction will register itself to the TransportService. this.updateSnapshotStatusHandler = new UpdateSnapshotStatusAction( diff --git a/server/src/test/java/org/opensearch/gateway/remote/model/RemotePinnedTimestampsTests.java b/server/src/test/java/org/opensearch/gateway/remote/model/RemotePinnedTimestampsTests.java new file mode 100644 index 0000000000000..309263a634265 --- /dev/null +++ b/server/src/test/java/org/opensearch/gateway/remote/model/RemotePinnedTimestampsTests.java @@ -0,0 +1,101 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.gateway.remote.model; + +import org.opensearch.common.compress.DeflateCompressor; +import org.opensearch.common.io.stream.BytesStreamOutput; +import org.opensearch.core.common.io.stream.BytesStreamInput; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.core.compress.Compressor; +import org.opensearch.test.OpenSearchTestCase; +import org.junit.Before; + +import java.io.IOException; +import java.io.InputStream; +import java.util.Arrays; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +public class RemotePinnedTimestampsTests extends OpenSearchTestCase { + + private RemotePinnedTimestamps remotePinnedTimestamps; + + @Before + public void setup() { + Compressor compressor = new DeflateCompressor(); + remotePinnedTimestamps = new RemotePinnedTimestamps("testClusterUUID", compressor); + } + + public void testGenerateBlobFileName() { + String fileName = remotePinnedTimestamps.generateBlobFileName(); + assertTrue(fileName.startsWith(RemotePinnedTimestamps.PINNED_TIMESTAMPS)); + assertEquals(fileName, remotePinnedTimestamps.getBlobFileName()); + } + + public void testSerializeAndDeserialize() throws IOException { + RemotePinnedTimestamps.PinnedTimestamps pinnedTimestamps = new RemotePinnedTimestamps.PinnedTimestamps(new HashMap<>()); + pinnedTimestamps.pin(1000L, "entity1"); + pinnedTimestamps.pin(2000L, "entity2"); + remotePinnedTimestamps.setPinnedTimestamps(pinnedTimestamps); + + InputStream serialized = remotePinnedTimestamps.serialize(); + RemotePinnedTimestamps.PinnedTimestamps deserialized = remotePinnedTimestamps.deserialize(serialized); + + assertEquals(pinnedTimestamps.getPinnedTimestampPinningEntityMap(), deserialized.getPinnedTimestampPinningEntityMap()); + } + + public void testSetAndGetPinnedTimestamps() { + RemotePinnedTimestamps.PinnedTimestamps pinnedTimestamps = new RemotePinnedTimestamps.PinnedTimestamps(new HashMap<>()); + remotePinnedTimestamps.setPinnedTimestamps(pinnedTimestamps); + assertEquals(pinnedTimestamps, remotePinnedTimestamps.getPinnedTimestamps()); + } + + public void testPinnedTimestampsPin() { + RemotePinnedTimestamps.PinnedTimestamps pinnedTimestamps = new RemotePinnedTimestamps.PinnedTimestamps(new HashMap<>()); + pinnedTimestamps.pin(1000L, "entity1"); + pinnedTimestamps.pin(1000L, "entity2"); + pinnedTimestamps.pin(2000L, "entity3"); + + Map> expected = new HashMap<>(); + expected.put(1000L, Arrays.asList("entity1", "entity2")); + expected.put(2000L, List.of("entity3")); + + assertEquals(expected, pinnedTimestamps.getPinnedTimestampPinningEntityMap()); + } + + public void testPinnedTimestampsUnpin() { + RemotePinnedTimestamps.PinnedTimestamps pinnedTimestamps = new RemotePinnedTimestamps.PinnedTimestamps(new HashMap<>()); + pinnedTimestamps.pin(1000L, "entity1"); + pinnedTimestamps.pin(1000L, "entity2"); + pinnedTimestamps.pin(2000L, "entity3"); + + pinnedTimestamps.unpin(1000L, "entity1"); + pinnedTimestamps.unpin(2000L, "entity3"); + + Map> expected = new HashMap<>(); + expected.put(1000L, List.of("entity2")); + + assertEquals(expected, pinnedTimestamps.getPinnedTimestampPinningEntityMap()); + } + + public void testPinnedTimestampsReadFromAndWriteTo() throws IOException { + RemotePinnedTimestamps.PinnedTimestamps original = new RemotePinnedTimestamps.PinnedTimestamps(new HashMap<>()); + original.pin(1000L, "entity1"); + original.pin(2000L, "entity2"); + + BytesStreamOutput out = new BytesStreamOutput(); + original.writeTo(out); + + StreamInput in = new BytesStreamInput(out.bytes().toBytesRef().bytes); + RemotePinnedTimestamps.PinnedTimestamps deserialized = RemotePinnedTimestamps.PinnedTimestamps.readFrom(in); + + assertEquals(original.getPinnedTimestampPinningEntityMap(), deserialized.getPinnedTimestampPinningEntityMap()); + } +} diff --git a/server/src/test/java/org/opensearch/snapshots/SnapshotResiliencyTests.java b/server/src/test/java/org/opensearch/snapshots/SnapshotResiliencyTests.java index 9c58fc8fde084..769dfeb37ff8d 100644 --- a/server/src/test/java/org/opensearch/snapshots/SnapshotResiliencyTests.java +++ b/server/src/test/java/org/opensearch/snapshots/SnapshotResiliencyTests.java @@ -2015,7 +2015,8 @@ public void onFailure(final Exception e) { indexNameExpressionResolver, repositoriesService, transportService, - actionFilters + actionFilters, + null ); nodeEnv = new NodeEnvironment(settings, environment); final NamedXContentRegistry namedXContentRegistry = new NamedXContentRegistry(Collections.emptyList()); From ef87b397e3c4d9571e6a29a96e438c3e006a0ea7 Mon Sep 17 00:00:00 2001 From: Finn Date: Thu, 15 Aug 2024 10:04:47 -0700 Subject: [PATCH 36/40] Remove filter rewrite optimization for range aggregations when segment is not effective match all (#15194) --------- Signed-off-by: Finn Carroll --- CHANGELOG.md | 1 + .../search.aggregation/360_date_histogram.yml | 91 ++++++++++++ .../test/search.aggregation/40_range.yml | 79 +++++++++++ .../filterrewrite/AggregatorBridge.java | 18 +++ .../DateHistogramAggregatorBridge.java | 17 --- .../filterrewrite/RangeAggregatorBridge.java | 1 - .../bucket/range/RangeAggregator.java | 6 +- .../bucket/range/RangeAggregatorTests.java | 129 +++++++++++++++++- 8 files changed, 320 insertions(+), 22 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1def0448ad24c..511dd77c24e22 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -49,6 +49,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), - Fix delete index template failed when the index template matches a data stream but is unused ([#15080](https://github.com/opensearch-project/OpenSearch/pull/15080)) - Fix array_index_out_of_bounds_exception when indexing documents with field name containing only dot ([#15126](https://github.com/opensearch-project/OpenSearch/pull/15126)) - Fixed array field name omission in flat_object function for nested JSON ([#13620](https://github.com/opensearch-project/OpenSearch/pull/13620)) +- Fix range aggregation optimization ignoring top level queries ([#15194](https://github.com/opensearch-project/OpenSearch/pull/15194)) ### Security diff --git a/rest-api-spec/src/main/resources/rest-api-spec/test/search.aggregation/360_date_histogram.yml b/rest-api-spec/src/main/resources/rest-api-spec/test/search.aggregation/360_date_histogram.yml index 0ea9d3de00926..8c8a98b2db22c 100644 --- a/rest-api-spec/src/main/resources/rest-api-spec/test/search.aggregation/360_date_histogram.yml +++ b/rest-api-spec/src/main/resources/rest-api-spec/test/search.aggregation/360_date_histogram.yml @@ -61,3 +61,94 @@ setup: - match: { aggregations.histo.buckets.8.doc_count: 1 } - match: { aggregations.histo.buckets.12.key_as_string: "2016-06-01T00:00:00.000Z" } - match: { aggregations.histo.buckets.12.doc_count: 1 } + +--- +"Date histogram aggregation w/ filter query test": + - skip: + version: " - 2.99.99" + reason: Backport fix to 2.16 + + - do: + bulk: + refresh: true + index: dhisto-agg-w-query + body: + - '{"index": {}}' + - '{"routing": "route1", "date": "2024-08-12", "dow": "monday"}' + - '{"index": {}}' + - '{"routing": "route1", "date": "2024-08-14", "dow": "wednesday"}' + - '{"index": {}}' + - '{"routing": "route1", "date": "2024-08-19", "dow": "monday"}' + - '{"index": {}}' + - '{"routing": "route2", "date": "2024-08-13", "dow": "tuesday"}' + - '{"index": {}}' + - '{"routing": "route2", "date": "2024-08-15", "dow": "thursday"}' + + - do: + search: + index: dhisto-agg-w-query + body: + query: + bool: + must: + match_all: {} + filter: + - terms: + routing: + - "route1" + aggregations: + weekHisto: + date_histogram: + field: date + calendar_interval: week + _source: false + + - match: { hits.total.value: 3 } + - match: { aggregations.weekHisto.buckets.0.doc_count: 2 } + - match: { aggregations.weekHisto.buckets.1.doc_count: 1 } + +--- +"Date histogram aggregation w/ shared field range test": + - do: + bulk: + refresh: true + index: dhisto-agg-w-query + body: + - '{"index": {}}' + - '{"date": "2024-10-31"}' + - '{"index": {}}' + - '{"date": "2024-11-11"}' + - '{"index": {}}' + - '{"date": "2024-11-28"}' + - '{"index": {}}' + - '{"date": "2024-12-25"}' + - '{"index": {}}' + - '{"date": "2025-01-01"}' + - '{"index": {}}' + - '{"date": "2025-02-14"}' + + - do: + search: + index: dhisto-agg-w-query + body: + profile: true + query: + range: + date: + gte: "2024-01-01" + lt: "2025-01-01" + aggregations: + monthHisto: + date_histogram: + field: date + calendar_interval: month + _source: false + + - match: { hits.total.value: 4 } + - match: { aggregations.monthHisto.buckets.0.doc_count: 1 } + - match: { aggregations.monthHisto.buckets.1.doc_count: 2 } + - match: { aggregations.monthHisto.buckets.2.doc_count: 1 } + - match: { profile.shards.0.aggregations.0.debug.optimized_segments: 1 } + - match: { profile.shards.0.aggregations.0.debug.unoptimized_segments: 0 } + - match: { profile.shards.0.aggregations.0.debug.leaf_visited: 0 } + - match: { profile.shards.0.aggregations.0.debug.inner_visited: 0 } diff --git a/rest-api-spec/src/main/resources/rest-api-spec/test/search.aggregation/40_range.yml b/rest-api-spec/src/main/resources/rest-api-spec/test/search.aggregation/40_range.yml index 80aad96ce1f6b..1e1d2b0706d6b 100644 --- a/rest-api-spec/src/main/resources/rest-api-spec/test/search.aggregation/40_range.yml +++ b/rest-api-spec/src/main/resources/rest-api-spec/test/search.aggregation/40_range.yml @@ -673,3 +673,82 @@ setup: - match: { aggregations.my_range.buckets.3.from: 1.5 } - is_false: aggregations.my_range.buckets.3.to - match: { aggregations.my_range.buckets.3.doc_count: 2 } + +--- +"Filter query w/ aggregation test": + - skip: + version: " - 2.99.99" + reason: Backport fix to 2.16 + + - do: + bulk: + refresh: true + index: range-agg-w-query + body: + - '{"index": {}}' + - '{"routing": "route1", "v": -10, "date": "2024-10-29"}' + - '{"index": {}}' + - '{"routing": "route1", "v": -5, "date": "2024-10-30"}' + - '{"index": {}}' + - '{"routing": "route1", "v": 10, "date": "2024-10-31"}' + - '{"index": {}}' + - '{"routing": "route2", "v": 15, "date": "2024-11-01"}' + - '{"index": {}}' + - '{"routing": "route2", "v": 20, "date": "2024-11-02"}' + + - do: + search: + index: range-agg-w-query + body: + query: + bool: + must: + match_all: {} + filter: + - terms: + routing: + - "route1" + aggregations: + NegPosAgg: + range: + field: v + keyed: true + ranges: + - to: 0 + key: "0" + - from: 0 + key: "1" + _source: false + + - match: { hits.total.value: 3 } + - match: { aggregations.NegPosAgg.buckets.0.doc_count: 2 } + - match: { aggregations.NegPosAgg.buckets.1.doc_count: 1 } + + - do: + search: + index: range-agg-w-query + body: + query: + bool: + must: + match_all: {} + filter: + - terms: + routing: + - "route1" + aggregations: + HalloweenAgg: + date_range: + field: date + format: "yyyy-MM-dd" + keyed: true + ranges: + - to: "2024-11-01" + key: "to-october" + - from: "2024-11-01" + key: "from-september" + _source: false + + - match: { hits.total.value: 3 } + - match: { aggregations.HalloweenAgg.buckets.to-october.doc_count: 3 } + - match: { aggregations.HalloweenAgg.buckets.from-september.doc_count: 0 } diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/AggregatorBridge.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/AggregatorBridge.java index 6b34582b259ea..145a60373b4f3 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/AggregatorBridge.java +++ b/server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/AggregatorBridge.java @@ -10,7 +10,10 @@ import org.apache.lucene.index.LeafReaderContext; import org.apache.lucene.index.PointValues; +import org.apache.lucene.search.ScoreMode; +import org.apache.lucene.search.Weight; import org.opensearch.index.mapper.MappedFieldType; +import org.opensearch.search.internal.SearchContext; import java.io.IOException; import java.util.function.BiConsumer; @@ -81,4 +84,19 @@ abstract FilterRewriteOptimizationContext.DebugInfo tryOptimize( BiConsumer incrementDocCount, Ranges ranges ) throws IOException; + + /** + * Checks whether the top level query matches all documents on the segment + * + *

      This method creates a weight from the search context's query and checks whether the weight's + * document count matches the total number of documents in the leaf reader context. + * + * @param ctx the search context + * @param leafCtx the leaf reader context for the segment + * @return {@code true} if the segment matches all documents, {@code false} otherwise + */ + public static boolean segmentMatchAll(SearchContext ctx, LeafReaderContext leafCtx) throws IOException { + Weight weight = ctx.query().rewrite(ctx.searcher()).createWeight(ctx.searcher(), ScoreMode.COMPLETE_NO_SCORES, 1f); + return weight != null && weight.count(leafCtx) == leafCtx.reader().numDocs(); + } } diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/DateHistogramAggregatorBridge.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/DateHistogramAggregatorBridge.java index 8bff3fdc5fefb..c780732a5ddce 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/DateHistogramAggregatorBridge.java +++ b/server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/DateHistogramAggregatorBridge.java @@ -11,8 +11,6 @@ import org.apache.lucene.document.LongPoint; import org.apache.lucene.index.LeafReaderContext; import org.apache.lucene.index.PointValues; -import org.apache.lucene.search.ScoreMode; -import org.apache.lucene.search.Weight; import org.opensearch.common.Rounding; import org.opensearch.index.mapper.DateFieldMapper; import org.opensearch.index.mapper.MappedFieldType; @@ -156,19 +154,4 @@ private static long getBucketOrd(long bucketOrd) { * Provides a function to produce bucket ordinals from the lower bound of the range */ protected abstract Function bucketOrdProducer(); - - /** - * Checks whether the top level query matches all documents on the segment - * - *

      This method creates a weight from the search context's query and checks whether the weight's - * document count matches the total number of documents in the leaf reader context. - * - * @param ctx the search context - * @param leafCtx the leaf reader context for the segment - * @return {@code true} if the segment matches all documents, {@code false} otherwise - */ - public static boolean segmentMatchAll(SearchContext ctx, LeafReaderContext leafCtx) throws IOException { - Weight weight = ctx.query().rewrite(ctx.searcher()).createWeight(ctx.searcher(), ScoreMode.COMPLETE_NO_SCORES, 1f); - return weight != null && weight.count(leafCtx) == leafCtx.reader().numDocs(); - } } diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/RangeAggregatorBridge.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/RangeAggregatorBridge.java index b590a444c8b04..fc1bcd83f2c1b 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/RangeAggregatorBridge.java +++ b/server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/RangeAggregatorBridge.java @@ -80,7 +80,6 @@ final FilterRewriteOptimizationContext.DebugInfo tryOptimize( Ranges ranges ) throws IOException { int size = Integer.MAX_VALUE; - BiConsumer incrementFunc = (activeIndex, docCount) -> { long bucketOrd = bucketOrdProducer().apply(activeIndex); incrementDocCount.accept(bucketOrd, (long) docCount); diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/range/RangeAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/range/RangeAggregator.java index 17461f228e993..1b6e0fe8c8d3f 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/bucket/range/RangeAggregator.java +++ b/server/src/main/java/org/opensearch/search/aggregations/bucket/range/RangeAggregator.java @@ -70,6 +70,7 @@ import java.util.function.Function; import static org.opensearch.core.xcontent.ConstructingObjectParser.optionalConstructorArg; +import static org.opensearch.search.aggregations.bucket.filterrewrite.AggregatorBridge.segmentMatchAll; /** * Aggregate all docs that match given ranges. @@ -310,8 +311,9 @@ public ScoreMode scoreMode() { @Override public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBucketCollector sub) throws IOException { - boolean optimized = filterRewriteOptimizationContext.tryOptimize(ctx, this::incrementBucketDocCount, false); - if (optimized) throw new CollectionTerminatedException(); + if (segmentMatchAll(context, ctx) && filterRewriteOptimizationContext.tryOptimize(ctx, this::incrementBucketDocCount, false)) { + throw new CollectionTerminatedException(); + } final SortedNumericDoubleValues values = valuesSource.doubleValues(ctx); return new LeafBucketCollectorBase(sub, values) { diff --git a/server/src/test/java/org/opensearch/search/aggregations/bucket/range/RangeAggregatorTests.java b/server/src/test/java/org/opensearch/search/aggregations/bucket/range/RangeAggregatorTests.java index 7e796b684e869..630326967aae1 100644 --- a/server/src/test/java/org/opensearch/search/aggregations/bucket/range/RangeAggregatorTests.java +++ b/server/src/test/java/org/opensearch/search/aggregations/bucket/range/RangeAggregatorTests.java @@ -32,6 +32,9 @@ package org.opensearch.search.aggregations.bucket.range; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.IntPoint; +import org.apache.lucene.document.KeywordField; import org.apache.lucene.document.NumericDocValuesField; import org.apache.lucene.document.SortedNumericDocValuesField; import org.apache.lucene.document.SortedSetDocValuesField; @@ -39,9 +42,13 @@ import org.apache.lucene.index.IndexReader; import org.apache.lucene.index.IndexWriter; import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.Term; +import org.apache.lucene.search.BooleanClause; +import org.apache.lucene.search.BooleanQuery; import org.apache.lucene.search.IndexSearcher; import org.apache.lucene.search.MatchAllDocsQuery; import org.apache.lucene.search.Query; +import org.apache.lucene.search.TermQuery; import org.apache.lucene.store.Directory; import org.apache.lucene.tests.index.RandomIndexWriter; import org.apache.lucene.tests.util.TestUtil; @@ -54,6 +61,7 @@ import org.opensearch.index.mapper.MappedFieldType; import org.opensearch.index.mapper.NumberFieldMapper.NumberFieldType; import org.opensearch.index.mapper.NumberFieldMapper.NumberType; +import org.opensearch.index.mapper.ParseContext.Document; import org.opensearch.search.aggregations.AggregationBuilder; import org.opensearch.search.aggregations.AggregatorTestCase; import org.opensearch.search.aggregations.CardinalityUpperBound; @@ -65,6 +73,7 @@ import java.io.IOException; import java.time.ZoneOffset; import java.time.ZonedDateTime; +import java.util.ArrayList; import java.util.Collections; import java.util.LinkedHashMap; import java.util.LinkedList; @@ -457,6 +466,29 @@ public void testFloatType() throws IOException { ); } + public void testTopLevelRangeQuery() throws IOException { + NumberFieldType fieldType = new NumberFieldType(NumberType.INTEGER.typeName(), NumberType.INTEGER); + String fieldName = fieldType.numberType().typeName(); + Query query = IntPoint.newRangeQuery(fieldName, 5, 20); + + testRewriteOptimizationCase( + fieldType, + new double[][] { { 0.0, 10.0 }, { 10.0, 20.0 } }, + query, + new Number[] { 0.1, 4.0, 9, 11, 12, 19 }, + range -> { + List ranges = range.getBuckets(); + assertEquals(2, ranges.size()); + assertEquals("0.0-10.0", ranges.get(0).getKeyAsString()); + assertEquals(1, ranges.get(0).getDocCount()); + assertEquals("10.0-20.0", ranges.get(1).getKeyAsString()); + assertEquals(3, ranges.get(1).getDocCount()); + assertTrue(AggregationInspectionHelper.hasValue(range)); + }, + false + ); + } + public void testUnsignedLongType() throws IOException { testRewriteOptimizationCase( new NumberFieldType(NumberType.UNSIGNED_LONG.typeName(), NumberType.UNSIGNED_LONG), @@ -493,6 +525,76 @@ public void testUnsignedLongType() throws IOException { ); } + /** + * Expect no optimization as top level query excludes documents. + */ + public void testTopLevelFilterTermQuery() throws IOException { + final String KEYWORD_FIELD_NAME = "route"; + final NumberFieldType NUM_FIELD_TYPE = new NumberFieldType(NumberType.DOUBLE.typeName(), NumberType.DOUBLE); + final NumberType numType = NUM_FIELD_TYPE.numberType(); + + BooleanQuery.Builder builder = new BooleanQuery.Builder(); + builder.setMinimumNumberShouldMatch(0); + builder.add(new TermQuery(new Term(KEYWORD_FIELD_NAME, "route1")), BooleanClause.Occur.MUST); + Query boolQuery = builder.build(); + + List docList = new ArrayList<>(); + for (int i = 0; i < 3; i++) + docList.add(new Document()); + + docList.get(0).addAll(numType.createFields(numType.typeName(), 3.0, true, true, false)); + docList.get(1).addAll(numType.createFields(numType.typeName(), 11.0, true, true, false)); + docList.get(2).addAll(numType.createFields(numType.typeName(), 15.0, true, true, false)); + docList.get(0).add(new KeywordField(KEYWORD_FIELD_NAME, "route1", Field.Store.NO)); + docList.get(1).add(new KeywordField(KEYWORD_FIELD_NAME, "route1", Field.Store.NO)); + docList.get(2).add(new KeywordField(KEYWORD_FIELD_NAME, "route2", Field.Store.NO)); + + testRewriteOptimizationCase(NUM_FIELD_TYPE, new double[][] { { 0.0, 10.0 }, { 10.0, 20.0 } }, boolQuery, docList, range -> { + List ranges = range.getBuckets(); + assertEquals(2, ranges.size()); + assertEquals("0.0-10.0", ranges.get(0).getKeyAsString()); + assertEquals(1, ranges.get(0).getDocCount()); + assertEquals("10.0-20.0", ranges.get(1).getKeyAsString()); + assertEquals(1, ranges.get(1).getDocCount()); + assertTrue(AggregationInspectionHelper.hasValue(range)); + }, false); + } + + /** + * Expect optimization as top level query is effective match all. + */ + public void testTopLevelEffectiveMatchAll() throws IOException { + final String KEYWORD_FIELD_NAME = "route"; + final NumberFieldType NUM_FIELD_TYPE = new NumberFieldType(NumberType.DOUBLE.typeName(), NumberType.DOUBLE); + final NumberType numType = NUM_FIELD_TYPE.numberType(); + + BooleanQuery.Builder builder = new BooleanQuery.Builder(); + builder.setMinimumNumberShouldMatch(0); + builder.add(new TermQuery(new Term(KEYWORD_FIELD_NAME, "route1")), BooleanClause.Occur.MUST); + Query boolQuery = builder.build(); + + List docList = new ArrayList<>(); + for (int i = 0; i < 3; i++) + docList.add(new Document()); + + docList.get(0).addAll(numType.createFields(numType.typeName(), 3.0, true, true, false)); + docList.get(1).addAll(numType.createFields(numType.typeName(), 11.0, true, true, false)); + docList.get(2).addAll(numType.createFields(numType.typeName(), 15.0, true, true, false)); + docList.get(0).add(new KeywordField(KEYWORD_FIELD_NAME, "route1", Field.Store.NO)); + docList.get(1).add(new KeywordField(KEYWORD_FIELD_NAME, "route1", Field.Store.NO)); + docList.get(2).add(new KeywordField(KEYWORD_FIELD_NAME, "route1", Field.Store.NO)); + + testRewriteOptimizationCase(NUM_FIELD_TYPE, new double[][] { { 0.0, 10.0 }, { 10.0, 20.0 } }, boolQuery, docList, range -> { + List ranges = range.getBuckets(); + assertEquals(2, ranges.size()); + assertEquals("0.0-10.0", ranges.get(0).getKeyAsString()); + assertEquals(1, ranges.get(0).getDocCount()); + assertEquals("10.0-20.0", ranges.get(1).getKeyAsString()); + assertEquals(2, ranges.get(1).getDocCount()); + assertTrue(AggregationInspectionHelper.hasValue(range)); + }, true); + } + private void testCase( Query query, CheckedConsumer buildIndex, @@ -556,11 +658,34 @@ private void testRewriteOptimizationCase( ) throws IOException { NumberType numberType = fieldType.numberType(); String fieldName = numberType.typeName(); + List docs = new ArrayList<>(); + + for (Number dataPoint : dataPoints) { + Document doc = new Document(); + List fieldList = numberType.createFields(fieldName, dataPoint, true, true, false); + for (Field fld : fieldList) + doc.add(fld); + docs.add(doc); + } + + testRewriteOptimizationCase(fieldType, ranges, query, docs, verify, optimized); + } + + private void testRewriteOptimizationCase( + NumberFieldType fieldType, + double[][] ranges, + Query query, + List documents, + Consumer> verify, + boolean optimized + ) throws IOException { + NumberType numberType = fieldType.numberType(); + String fieldName = numberType.typeName(); try (Directory directory = newDirectory()) { try (IndexWriter indexWriter = new IndexWriter(directory, new IndexWriterConfig().setCodec(TestUtil.getDefaultCodec()))) { - for (Number dataPoint : dataPoints) { - indexWriter.addDocument(numberType.createFields(fieldName, dataPoint, true, true, false)); + for (Document doc : documents) { + indexWriter.addDocument(doc); } } From b31627986fde9e9a8ed0cd1e94529df965acc5b1 Mon Sep 17 00:00:00 2001 From: Andriy Redko Date: Thu, 15 Aug 2024 19:40:11 -0400 Subject: [PATCH 37/40] Fix Gradle runtime JDK version post 8.10 version upgrade (#15270) Signed-off-by: Andriy Redko --- .ci/java-versions.properties | 3 ++- buildSrc/src/main/resources/minimumCompilerVersion | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/.ci/java-versions.properties b/.ci/java-versions.properties index f73122ee21a6b..e290bda773f68 100644 --- a/.ci/java-versions.properties +++ b/.ci/java-versions.properties @@ -13,7 +13,8 @@ # build and test OpenSearch for this branch. Valid Java versions # are 'java' or 'openjdk' followed by the major release number. -OPENSEARCH_BUILD_JAVA=openjdk11 +# See please https://docs.gradle.org/8.10/userguide/upgrading_version_8.html#minimum_daemon_jvm_version +OPENSEARCH_BUILD_JAVA=openjdk17 OPENSEARCH_RUNTIME_JAVA=java11 GRADLE_TASK=build GRADLE_EXTRA_ARGS= diff --git a/buildSrc/src/main/resources/minimumCompilerVersion b/buildSrc/src/main/resources/minimumCompilerVersion index 8351c19397f4f..98d9bcb75a685 100644 --- a/buildSrc/src/main/resources/minimumCompilerVersion +++ b/buildSrc/src/main/resources/minimumCompilerVersion @@ -1 +1 @@ -14 +17 From cbe7921b0af4d6ad721239ebd3ce2ca4f807d437 Mon Sep 17 00:00:00 2001 From: Andriy Redko Date: Fri, 16 Aug 2024 11:30:19 -0400 Subject: [PATCH 38/40] Bump Project Reactor to 3.5.20 and Reactor Netty to 1.1.22 (#15262) Signed-off-by: Andriy Redko --- CHANGELOG.md | 2 ++ buildSrc/version.properties | 4 ++-- client/rest/licenses/reactor-core-3.5.19.jar.sha1 | 1 - client/rest/licenses/reactor-core-3.5.20.jar.sha1 | 1 + .../licenses/reactor-netty-core-1.1.21.jar.sha1 | 1 - .../licenses/reactor-netty-core-1.1.22.jar.sha1 | 1 + .../licenses/reactor-netty-http-1.1.21.jar.sha1 | 1 - .../licenses/reactor-netty-http-1.1.22.jar.sha1 | 1 + .../licenses/reactor-netty-core-1.1.21.jar.sha1 | 1 - .../licenses/reactor-netty-core-1.1.22.jar.sha1 | 1 + .../licenses/reactor-netty-http-1.1.21.jar.sha1 | 1 - .../licenses/reactor-netty-http-1.1.22.jar.sha1 | 1 + server/licenses/reactor-core-3.5.19.jar.sha1 | 1 - server/licenses/reactor-core-3.5.20.jar.sha1 | 1 + 14 files changed, 10 insertions(+), 8 deletions(-) delete mode 100644 client/rest/licenses/reactor-core-3.5.19.jar.sha1 create mode 100644 client/rest/licenses/reactor-core-3.5.20.jar.sha1 delete mode 100644 plugins/repository-azure/licenses/reactor-netty-core-1.1.21.jar.sha1 create mode 100644 plugins/repository-azure/licenses/reactor-netty-core-1.1.22.jar.sha1 delete mode 100644 plugins/repository-azure/licenses/reactor-netty-http-1.1.21.jar.sha1 create mode 100644 plugins/repository-azure/licenses/reactor-netty-http-1.1.22.jar.sha1 delete mode 100644 plugins/transport-reactor-netty4/licenses/reactor-netty-core-1.1.21.jar.sha1 create mode 100644 plugins/transport-reactor-netty4/licenses/reactor-netty-core-1.1.22.jar.sha1 delete mode 100644 plugins/transport-reactor-netty4/licenses/reactor-netty-http-1.1.21.jar.sha1 create mode 100644 plugins/transport-reactor-netty4/licenses/reactor-netty-http-1.1.22.jar.sha1 delete mode 100644 server/licenses/reactor-core-3.5.19.jar.sha1 create mode 100644 server/licenses/reactor-core-3.5.20.jar.sha1 diff --git a/CHANGELOG.md b/CHANGELOG.md index 511dd77c24e22..6cf1de210f98f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -33,6 +33,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), - Bump `com.azure:azure-core` from 1.49.1 to 1.51.0 ([#15111](https://github.com/opensearch-project/OpenSearch/pull/15111)) - Bump `org.xerial.snappy:snappy-java` from 1.1.10.5 to 1.1.10.6 ([#15207](https://github.com/opensearch-project/OpenSearch/pull/15207)) - Bump `com.azure:azure-xml` from 1.0.0 to 1.1.0 ([#15206](https://github.com/opensearch-project/OpenSearch/pull/15206)) +- Bump `reactor` from 3.5.19 to 3.5.20 ([#15262](https://github.com/opensearch-project/OpenSearch/pull/15262)) +- Bump `reactor-netty` from 1.1.21 to 1.1.22 ([#15262](https://github.com/opensearch-project/OpenSearch/pull/15262)) ### Changed - Add lower limit for primary and replica batch allocators timeout ([#14979](https://github.com/opensearch-project/OpenSearch/pull/14979)) diff --git a/buildSrc/version.properties b/buildSrc/version.properties index 08c45ef058716..9d7bbf6f8f769 100644 --- a/buildSrc/version.properties +++ b/buildSrc/version.properties @@ -33,8 +33,8 @@ netty = 4.1.112.Final joda = 2.12.7 # project reactor -reactor_netty = 1.1.21 -reactor = 3.5.19 +reactor_netty = 1.1.22 +reactor = 3.5.20 # client dependencies httpclient5 = 5.2.3 diff --git a/client/rest/licenses/reactor-core-3.5.19.jar.sha1 b/client/rest/licenses/reactor-core-3.5.19.jar.sha1 deleted file mode 100644 index 04b59d2faae04..0000000000000 --- a/client/rest/licenses/reactor-core-3.5.19.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -1d49ce1d0df79f28d3927da5f4c46a895b94335f \ No newline at end of file diff --git a/client/rest/licenses/reactor-core-3.5.20.jar.sha1 b/client/rest/licenses/reactor-core-3.5.20.jar.sha1 new file mode 100644 index 0000000000000..0c80be89f66c8 --- /dev/null +++ b/client/rest/licenses/reactor-core-3.5.20.jar.sha1 @@ -0,0 +1 @@ +1fc0f91e2b93778a974339d2c24363d7f34f90b4 \ No newline at end of file diff --git a/plugins/repository-azure/licenses/reactor-netty-core-1.1.21.jar.sha1 b/plugins/repository-azure/licenses/reactor-netty-core-1.1.21.jar.sha1 deleted file mode 100644 index 21c16c7430016..0000000000000 --- a/plugins/repository-azure/licenses/reactor-netty-core-1.1.21.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -acb98bd08107287c454ce74e7b1ed8e7a018a662 \ No newline at end of file diff --git a/plugins/repository-azure/licenses/reactor-netty-core-1.1.22.jar.sha1 b/plugins/repository-azure/licenses/reactor-netty-core-1.1.22.jar.sha1 new file mode 100644 index 0000000000000..cc894568c5760 --- /dev/null +++ b/plugins/repository-azure/licenses/reactor-netty-core-1.1.22.jar.sha1 @@ -0,0 +1 @@ +08356b59b29f86e7142c9daca0434653a64ae64b \ No newline at end of file diff --git a/plugins/repository-azure/licenses/reactor-netty-http-1.1.21.jar.sha1 b/plugins/repository-azure/licenses/reactor-netty-http-1.1.21.jar.sha1 deleted file mode 100644 index 648df22873d56..0000000000000 --- a/plugins/repository-azure/licenses/reactor-netty-http-1.1.21.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -b83542bb35630ef815b4177e3c670f62e952e695 \ No newline at end of file diff --git a/plugins/repository-azure/licenses/reactor-netty-http-1.1.22.jar.sha1 b/plugins/repository-azure/licenses/reactor-netty-http-1.1.22.jar.sha1 new file mode 100644 index 0000000000000..2402813f831ce --- /dev/null +++ b/plugins/repository-azure/licenses/reactor-netty-http-1.1.22.jar.sha1 @@ -0,0 +1 @@ +2faf64b3822b0512f15d72a325e2826eb8564413 \ No newline at end of file diff --git a/plugins/transport-reactor-netty4/licenses/reactor-netty-core-1.1.21.jar.sha1 b/plugins/transport-reactor-netty4/licenses/reactor-netty-core-1.1.21.jar.sha1 deleted file mode 100644 index 21c16c7430016..0000000000000 --- a/plugins/transport-reactor-netty4/licenses/reactor-netty-core-1.1.21.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -acb98bd08107287c454ce74e7b1ed8e7a018a662 \ No newline at end of file diff --git a/plugins/transport-reactor-netty4/licenses/reactor-netty-core-1.1.22.jar.sha1 b/plugins/transport-reactor-netty4/licenses/reactor-netty-core-1.1.22.jar.sha1 new file mode 100644 index 0000000000000..cc894568c5760 --- /dev/null +++ b/plugins/transport-reactor-netty4/licenses/reactor-netty-core-1.1.22.jar.sha1 @@ -0,0 +1 @@ +08356b59b29f86e7142c9daca0434653a64ae64b \ No newline at end of file diff --git a/plugins/transport-reactor-netty4/licenses/reactor-netty-http-1.1.21.jar.sha1 b/plugins/transport-reactor-netty4/licenses/reactor-netty-http-1.1.21.jar.sha1 deleted file mode 100644 index 648df22873d56..0000000000000 --- a/plugins/transport-reactor-netty4/licenses/reactor-netty-http-1.1.21.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -b83542bb35630ef815b4177e3c670f62e952e695 \ No newline at end of file diff --git a/plugins/transport-reactor-netty4/licenses/reactor-netty-http-1.1.22.jar.sha1 b/plugins/transport-reactor-netty4/licenses/reactor-netty-http-1.1.22.jar.sha1 new file mode 100644 index 0000000000000..2402813f831ce --- /dev/null +++ b/plugins/transport-reactor-netty4/licenses/reactor-netty-http-1.1.22.jar.sha1 @@ -0,0 +1 @@ +2faf64b3822b0512f15d72a325e2826eb8564413 \ No newline at end of file diff --git a/server/licenses/reactor-core-3.5.19.jar.sha1 b/server/licenses/reactor-core-3.5.19.jar.sha1 deleted file mode 100644 index 04b59d2faae04..0000000000000 --- a/server/licenses/reactor-core-3.5.19.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -1d49ce1d0df79f28d3927da5f4c46a895b94335f \ No newline at end of file diff --git a/server/licenses/reactor-core-3.5.20.jar.sha1 b/server/licenses/reactor-core-3.5.20.jar.sha1 new file mode 100644 index 0000000000000..0c80be89f66c8 --- /dev/null +++ b/server/licenses/reactor-core-3.5.20.jar.sha1 @@ -0,0 +1 @@ +1fc0f91e2b93778a974339d2c24363d7f34f90b4 \ No newline at end of file From a900a16ce8c076488324993ec963f6a6f0e9e269 Mon Sep 17 00:00:00 2001 From: Ruirui Zhang Date: Fri, 16 Aug 2024 10:59:07 -0700 Subject: [PATCH 39/40] Add Get QueryGroup API Logic (#14709) * Add Get QueryGroup API Logic Signed-off-by: Ruirui Zhang * add to changelog Signed-off-by: Ruirui Zhang * fix javadoc Signed-off-by: Ruirui Zhang * change GetQueryGroupAction NAME and add more tests Signed-off-by: Ruirui Zhang * add more unit tests Signed-off-by: Ruirui Zhang * fix spotlessapply Signed-off-by: Ruirui Zhang * addressed comments Signed-off-by: Ruirui Zhang * incorperate comments from create api PR Signed-off-by: Ruirui Zhang * use clustermanager to get the most recent querygroups Signed-off-by: Ruirui Zhang * address comments Signed-off-by: Ruirui Zhang * rebase with main Signed-off-by: Ruirui Zhang * add IT Signed-off-by: Ruirui Zhang * address comments Signed-off-by: Ruirui Zhang * fix IT Signed-off-by: Ruirui Zhang --- CHANGELOG.md | 1 + .../plugin/wlm/WorkloadManagementPlugin.java | 10 +- .../wlm/action/GetQueryGroupAction.java | 36 +++++ .../wlm/action/GetQueryGroupRequest.java | 64 ++++++++ .../wlm/action/GetQueryGroupResponse.java | 82 ++++++++++ .../action/TransportGetQueryGroupAction.java | 98 ++++++++++++ .../wlm/rest/RestGetQueryGroupAction.java | 68 ++++++++ .../service/QueryGroupPersistenceService.java | 20 +++ .../plugin/wlm/QueryGroupTestUtils.java | 7 +- .../wlm/action/GetQueryGroupRequestTests.java | 53 ++++++ .../action/GetQueryGroupResponseTests.java | 151 ++++++++++++++++++ .../TransportGetQueryGroupActionTests.java | 47 ++++++ .../QueryGroupPersistenceServiceTests.java | 54 +++++++ .../api/get_query_group_context.json | 25 +++ ...ate_query_group.yml => 10_query_group.yml} | 20 ++- .../cluster/metadata/QueryGroup.java | 11 +- 16 files changed, 738 insertions(+), 9 deletions(-) create mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/GetQueryGroupAction.java create mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/GetQueryGroupRequest.java create mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/GetQueryGroupResponse.java create mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportGetQueryGroupAction.java create mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestGetQueryGroupAction.java create mode 100644 plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/GetQueryGroupRequestTests.java create mode 100644 plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/GetQueryGroupResponseTests.java create mode 100644 plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/TransportGetQueryGroupActionTests.java create mode 100644 plugins/workload-management/src/yamlRestTest/resources/rest-api-spec/api/get_query_group_context.json rename plugins/workload-management/src/yamlRestTest/resources/rest-api-spec/test/wlm/{10_create_query_group.yml => 10_query_group.yml} (77%) diff --git a/CHANGELOG.md b/CHANGELOG.md index 6cf1de210f98f..9fe700d7190fa 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -17,6 +17,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), - [Concurrent Segment Search] Support composite aggregations with scripting ([#15072](https://github.com/opensearch-project/OpenSearch/pull/15072)) - Add `rangeQuery` and `regexpQuery` for `constant_keyword` field type ([#14711](https://github.com/opensearch-project/OpenSearch/pull/14711)) - Add took time to request nodes stats ([#15054](https://github.com/opensearch-project/OpenSearch/pull/15054)) +- [Workload Management] Add Get QueryGroup API Logic ([14709](https://github.com/opensearch-project/OpenSearch/pull/14709)) - [Workload Management] QueryGroup resource tracking framework changes ([#13897](https://github.com/opensearch-project/OpenSearch/pull/13897)) - Add slice execution listeners to SearchOperationListener interface ([#15153](https://github.com/opensearch-project/OpenSearch/pull/15153)) diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java index 80807f0d5bc37..6b4496af76dc3 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java @@ -18,8 +18,11 @@ import org.opensearch.common.settings.SettingsFilter; import org.opensearch.core.action.ActionResponse; import org.opensearch.plugin.wlm.action.CreateQueryGroupAction; +import org.opensearch.plugin.wlm.action.GetQueryGroupAction; import org.opensearch.plugin.wlm.action.TransportCreateQueryGroupAction; +import org.opensearch.plugin.wlm.action.TransportGetQueryGroupAction; import org.opensearch.plugin.wlm.rest.RestCreateQueryGroupAction; +import org.opensearch.plugin.wlm.rest.RestGetQueryGroupAction; import org.opensearch.plugin.wlm.service.QueryGroupPersistenceService; import org.opensearch.plugins.ActionPlugin; import org.opensearch.plugins.Plugin; @@ -41,7 +44,10 @@ public WorkloadManagementPlugin() {} @Override public List> getActions() { - return List.of(new ActionPlugin.ActionHandler<>(CreateQueryGroupAction.INSTANCE, TransportCreateQueryGroupAction.class)); + return List.of( + new ActionPlugin.ActionHandler<>(CreateQueryGroupAction.INSTANCE, TransportCreateQueryGroupAction.class), + new ActionPlugin.ActionHandler<>(GetQueryGroupAction.INSTANCE, TransportGetQueryGroupAction.class) + ); } @Override @@ -54,7 +60,7 @@ public List getRestHandlers( IndexNameExpressionResolver indexNameExpressionResolver, Supplier nodesInCluster ) { - return List.of(new RestCreateQueryGroupAction()); + return List.of(new RestCreateQueryGroupAction(), new RestGetQueryGroupAction()); } @Override diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/GetQueryGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/GetQueryGroupAction.java new file mode 100644 index 0000000000000..0200185580f7d --- /dev/null +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/GetQueryGroupAction.java @@ -0,0 +1,36 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.plugin.wlm.action; + +import org.opensearch.action.ActionType; + +/** + * Transport action to get QueryGroup + * + * @opensearch.experimental + */ +public class GetQueryGroupAction extends ActionType { + + /** + * An instance of GetQueryGroupAction + */ + public static final GetQueryGroupAction INSTANCE = new GetQueryGroupAction(); + + /** + * Name for GetQueryGroupAction + */ + public static final String NAME = "cluster:admin/opensearch/wlm/query_group/_get"; + + /** + * Default constructor + */ + private GetQueryGroupAction() { + super(NAME, GetQueryGroupResponse::new); + } +} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/GetQueryGroupRequest.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/GetQueryGroupRequest.java new file mode 100644 index 0000000000000..0524c615a84e7 --- /dev/null +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/GetQueryGroupRequest.java @@ -0,0 +1,64 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.plugin.wlm.action; + +import org.opensearch.action.ActionRequestValidationException; +import org.opensearch.action.support.clustermanager.ClusterManagerNodeReadRequest; +import org.opensearch.cluster.metadata.QueryGroup; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.core.common.io.stream.StreamOutput; + +import java.io.IOException; + +/** + * Request for get QueryGroup + * + * @opensearch.experimental + */ +public class GetQueryGroupRequest extends ClusterManagerNodeReadRequest { + final String name; + + /** + * Default constructor for GetQueryGroupRequest + * @param name - name for the QueryGroup to get + */ + public GetQueryGroupRequest(String name) { + this.name = name; + } + + /** + * Constructor for GetQueryGroupRequest + * @param in - A {@link StreamInput} object + */ + public GetQueryGroupRequest(StreamInput in) throws IOException { + super(in); + name = in.readOptionalString(); + } + + @Override + public ActionRequestValidationException validate() { + if (name != null) { + QueryGroup.validateName(name); + } + return null; + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + super.writeTo(out); + out.writeOptionalString(name); + } + + /** + * Name getter + */ + public String getName() { + return name; + } +} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/GetQueryGroupResponse.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/GetQueryGroupResponse.java new file mode 100644 index 0000000000000..547c501e6a28e --- /dev/null +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/GetQueryGroupResponse.java @@ -0,0 +1,82 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.plugin.wlm.action; + +import org.opensearch.cluster.metadata.QueryGroup; +import org.opensearch.core.action.ActionResponse; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.core.common.io.stream.StreamOutput; +import org.opensearch.core.rest.RestStatus; +import org.opensearch.core.xcontent.ToXContent; +import org.opensearch.core.xcontent.ToXContentObject; +import org.opensearch.core.xcontent.XContentBuilder; + +import java.io.IOException; +import java.util.Collection; + +/** + * Response for the get API for QueryGroup + * + * @opensearch.experimental + */ +public class GetQueryGroupResponse extends ActionResponse implements ToXContent, ToXContentObject { + private final Collection queryGroups; + private final RestStatus restStatus; + + /** + * Constructor for GetQueryGroupResponse + * @param queryGroups - The QueryGroup list to be fetched + * @param restStatus - The rest status of the request + */ + public GetQueryGroupResponse(final Collection queryGroups, RestStatus restStatus) { + this.queryGroups = queryGroups; + this.restStatus = restStatus; + } + + /** + * Constructor for GetQueryGroupResponse + * @param in - A {@link StreamInput} object + */ + public GetQueryGroupResponse(StreamInput in) throws IOException { + this.queryGroups = in.readList(QueryGroup::new); + restStatus = RestStatus.readFrom(in); + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + out.writeCollection(queryGroups); + RestStatus.writeTo(out, restStatus); + } + + @Override + public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { + builder.startObject(); + builder.startArray("query_groups"); + for (QueryGroup group : queryGroups) { + group.toXContent(builder, params); + } + builder.endArray(); + builder.endObject(); + return builder; + } + + /** + * queryGroups getter + */ + public Collection getQueryGroups() { + return queryGroups; + } + + /** + * restStatus getter + */ + public RestStatus getRestStatus() { + return restStatus; + } +} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportGetQueryGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportGetQueryGroupAction.java new file mode 100644 index 0000000000000..51bb21b255511 --- /dev/null +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportGetQueryGroupAction.java @@ -0,0 +1,98 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.plugin.wlm.action; + +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.opensearch.ResourceNotFoundException; +import org.opensearch.action.support.ActionFilters; +import org.opensearch.action.support.clustermanager.TransportClusterManagerNodeReadAction; +import org.opensearch.cluster.ClusterState; +import org.opensearch.cluster.block.ClusterBlockException; +import org.opensearch.cluster.block.ClusterBlockLevel; +import org.opensearch.cluster.metadata.IndexNameExpressionResolver; +import org.opensearch.cluster.metadata.QueryGroup; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.inject.Inject; +import org.opensearch.core.action.ActionListener; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.core.rest.RestStatus; +import org.opensearch.plugin.wlm.service.QueryGroupPersistenceService; +import org.opensearch.search.pipeline.SearchPipelineService; +import org.opensearch.threadpool.ThreadPool; +import org.opensearch.transport.TransportService; + +import java.io.IOException; +import java.util.Collection; + +/** + * Transport action to get QueryGroup + * + * @opensearch.experimental + */ +public class TransportGetQueryGroupAction extends TransportClusterManagerNodeReadAction { + private static final Logger logger = LogManager.getLogger(SearchPipelineService.class); + + /** + * Constructor for TransportGetQueryGroupAction + * + * @param clusterService - a {@link ClusterService} object + * @param transportService - a {@link TransportService} object + * @param actionFilters - a {@link ActionFilters} object + * @param threadPool - a {@link ThreadPool} object + * @param indexNameExpressionResolver - a {@link IndexNameExpressionResolver} object + */ + @Inject + public TransportGetQueryGroupAction( + ClusterService clusterService, + TransportService transportService, + ActionFilters actionFilters, + ThreadPool threadPool, + IndexNameExpressionResolver indexNameExpressionResolver + ) { + super( + GetQueryGroupAction.NAME, + transportService, + clusterService, + threadPool, + actionFilters, + GetQueryGroupRequest::new, + indexNameExpressionResolver, + true + ); + } + + @Override + protected String executor() { + return ThreadPool.Names.SAME; + } + + @Override + protected GetQueryGroupResponse read(StreamInput in) throws IOException { + return new GetQueryGroupResponse(in); + } + + @Override + protected ClusterBlockException checkBlock(GetQueryGroupRequest request, ClusterState state) { + return state.blocks().globalBlockedException(ClusterBlockLevel.METADATA_READ); + } + + @Override + protected void clusterManagerOperation(GetQueryGroupRequest request, ClusterState state, ActionListener listener) + throws Exception { + final String name = request.getName(); + final Collection resultGroups = QueryGroupPersistenceService.getFromClusterStateMetadata(name, state); + + if (resultGroups.isEmpty() && name != null && !name.isEmpty()) { + logger.warn("No QueryGroup exists with the provided name: {}", name); + throw new ResourceNotFoundException("No QueryGroup exists with the provided name: " + name); + } + listener.onResponse(new GetQueryGroupResponse(resultGroups, RestStatus.OK)); + } +} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestGetQueryGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestGetQueryGroupAction.java new file mode 100644 index 0000000000000..c250bd2979e98 --- /dev/null +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestGetQueryGroupAction.java @@ -0,0 +1,68 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.plugin.wlm.rest; + +import org.opensearch.client.node.NodeClient; +import org.opensearch.core.rest.RestStatus; +import org.opensearch.core.xcontent.ToXContent; +import org.opensearch.plugin.wlm.action.GetQueryGroupAction; +import org.opensearch.plugin.wlm.action.GetQueryGroupRequest; +import org.opensearch.plugin.wlm.action.GetQueryGroupResponse; +import org.opensearch.rest.BaseRestHandler; +import org.opensearch.rest.BytesRestResponse; +import org.opensearch.rest.RestChannel; +import org.opensearch.rest.RestRequest; +import org.opensearch.rest.RestResponse; +import org.opensearch.rest.action.RestResponseListener; + +import java.io.IOException; +import java.util.List; + +import static org.opensearch.rest.RestRequest.Method.GET; + +/** + * Rest action to get a QueryGroup0 + * + * @opensearch.experimental + */ +public class RestGetQueryGroupAction extends BaseRestHandler { + + /** + * Constructor for RestGetQueryGroupAction + */ + public RestGetQueryGroupAction() {} + + @Override + public String getName() { + return "get_query_group"; + } + + /** + * The list of {@link Route}s that this RestHandler is responsible for handling. + */ + @Override + public List routes() { + return List.of(new Route(GET, "_wlm/query_group/{name}"), new Route(GET, "_wlm/query_group/")); + } + + @Override + protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException { + final GetQueryGroupRequest getQueryGroupRequest = new GetQueryGroupRequest(request.param("name")); + return channel -> client.execute(GetQueryGroupAction.INSTANCE, getQueryGroupRequest, getQueryGroupResponse(channel)); + } + + private RestResponseListener getQueryGroupResponse(final RestChannel channel) { + return new RestResponseListener<>(channel) { + @Override + public RestResponse buildResponse(final GetQueryGroupResponse response) throws Exception { + return new BytesRestResponse(RestStatus.OK, response.toXContent(channel.newBuilder(), ToXContent.EMPTY_PARAMS)); + } + }; + } +} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceService.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceService.java index b2164df561bf9..fe7080da78bbe 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceService.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceService.java @@ -26,9 +26,11 @@ import org.opensearch.plugin.wlm.action.CreateQueryGroupResponse; import org.opensearch.search.ResourceType; +import java.util.Collection; import java.util.EnumMap; import java.util.Map; import java.util.Optional; +import java.util.stream.Collectors; /** * This class defines the functions for QueryGroup persistence @@ -192,6 +194,24 @@ private Map calculateTotalUsage(Map ex return map; } + /** + * Get the QueryGroups with the specified name from cluster state + * @param name - the QueryGroup name we are getting + * @param currentState - current cluster state + */ + public static Collection getFromClusterStateMetadata(String name, ClusterState currentState) { + final Map currentGroups = currentState.getMetadata().queryGroups(); + if (name == null || name.isEmpty()) { + return currentGroups.values(); + } + return currentGroups.values() + .stream() + .filter(group -> group.getName().equals(name)) + .findAny() + .stream() + .collect(Collectors.toList()); + } + /** * maxQueryGroupCount getter */ diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/QueryGroupTestUtils.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/QueryGroupTestUtils.java index fc324853d9b34..5ba1ad5334712 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/QueryGroupTestUtils.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/QueryGroupTestUtils.java @@ -23,6 +23,7 @@ import org.opensearch.threadpool.ThreadPool; import java.util.ArrayList; +import java.util.Collection; import java.util.Comparator; import java.util.HashSet; import java.util.List; @@ -130,8 +131,10 @@ public static Tuple preparePersisten return new Tuple(queryGroupPersistenceService, clusterState); } - public static void assertEqualQueryGroups(List listOne, List listTwo) { - assertEquals(listOne.size(), listTwo.size()); + public static void assertEqualQueryGroups(Collection collectionOne, Collection collectionTwo) { + assertEquals(collectionOne.size(), collectionTwo.size()); + List listOne = new ArrayList<>(collectionOne); + List listTwo = new ArrayList<>(collectionTwo); listOne.sort(Comparator.comparing(QueryGroup::getName)); listTwo.sort(Comparator.comparing(QueryGroup::getName)); for (int i = 0; i < listOne.size(); i++) { diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/GetQueryGroupRequestTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/GetQueryGroupRequestTests.java new file mode 100644 index 0000000000000..32b5f7ec9e2c3 --- /dev/null +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/GetQueryGroupRequestTests.java @@ -0,0 +1,53 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.plugin.wlm.action; + +import org.opensearch.common.io.stream.BytesStreamOutput; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.plugin.wlm.QueryGroupTestUtils; +import org.opensearch.test.OpenSearchTestCase; + +import java.io.IOException; + +public class GetQueryGroupRequestTests extends OpenSearchTestCase { + + /** + * Test case to verify the serialization and deserialization of GetQueryGroupRequest. + */ + public void testSerialization() throws IOException { + GetQueryGroupRequest request = new GetQueryGroupRequest(QueryGroupTestUtils.NAME_ONE); + assertEquals(QueryGroupTestUtils.NAME_ONE, request.getName()); + BytesStreamOutput out = new BytesStreamOutput(); + request.writeTo(out); + StreamInput streamInput = out.bytes().streamInput(); + GetQueryGroupRequest otherRequest = new GetQueryGroupRequest(streamInput); + assertEquals(request.getName(), otherRequest.getName()); + } + + /** + * Test case to verify the serialization and deserialization of GetQueryGroupRequest when name is null. + */ + public void testSerializationWithNull() throws IOException { + GetQueryGroupRequest request = new GetQueryGroupRequest((String) null); + assertNull(request.getName()); + BytesStreamOutput out = new BytesStreamOutput(); + request.writeTo(out); + StreamInput streamInput = out.bytes().streamInput(); + GetQueryGroupRequest otherRequest = new GetQueryGroupRequest(streamInput); + assertEquals(request.getName(), otherRequest.getName()); + } + + /** + * Test case the validation function of GetQueryGroupRequest + */ + public void testValidation() { + GetQueryGroupRequest request = new GetQueryGroupRequest("a".repeat(51)); + assertThrows(IllegalArgumentException.class, request::validate); + } +} diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/GetQueryGroupResponseTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/GetQueryGroupResponseTests.java new file mode 100644 index 0000000000000..774f4b2d8db52 --- /dev/null +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/GetQueryGroupResponseTests.java @@ -0,0 +1,151 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.plugin.wlm.action; + +import org.opensearch.cluster.metadata.QueryGroup; +import org.opensearch.common.io.stream.BytesStreamOutput; +import org.opensearch.common.xcontent.json.JsonXContent; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.core.rest.RestStatus; +import org.opensearch.core.xcontent.ToXContent; +import org.opensearch.core.xcontent.XContentBuilder; +import org.opensearch.plugin.wlm.QueryGroupTestUtils; +import org.opensearch.test.OpenSearchTestCase; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; + +import static org.mockito.Mockito.mock; + +public class GetQueryGroupResponseTests extends OpenSearchTestCase { + + /** + * Test case to verify the serialization and deserialization of GetQueryGroupResponse. + */ + public void testSerializationSingleQueryGroup() throws IOException { + List list = new ArrayList<>(); + list.add(QueryGroupTestUtils.queryGroupOne); + GetQueryGroupResponse response = new GetQueryGroupResponse(list, RestStatus.OK); + assertEquals(response.getQueryGroups(), list); + + BytesStreamOutput out = new BytesStreamOutput(); + response.writeTo(out); + StreamInput streamInput = out.bytes().streamInput(); + + GetQueryGroupResponse otherResponse = new GetQueryGroupResponse(streamInput); + assertEquals(response.getRestStatus(), otherResponse.getRestStatus()); + QueryGroupTestUtils.assertEqualQueryGroups(response.getQueryGroups(), otherResponse.getQueryGroups()); + } + + /** + * Test case to verify the serialization and deserialization of GetQueryGroupResponse when the result contains multiple QueryGroups. + */ + public void testSerializationMultipleQueryGroup() throws IOException { + GetQueryGroupResponse response = new GetQueryGroupResponse(QueryGroupTestUtils.queryGroupList(), RestStatus.OK); + assertEquals(response.getQueryGroups(), QueryGroupTestUtils.queryGroupList()); + + BytesStreamOutput out = new BytesStreamOutput(); + response.writeTo(out); + StreamInput streamInput = out.bytes().streamInput(); + + GetQueryGroupResponse otherResponse = new GetQueryGroupResponse(streamInput); + assertEquals(response.getRestStatus(), otherResponse.getRestStatus()); + assertEquals(2, otherResponse.getQueryGroups().size()); + QueryGroupTestUtils.assertEqualQueryGroups(response.getQueryGroups(), otherResponse.getQueryGroups()); + } + + /** + * Test case to verify the serialization and deserialization of GetQueryGroupResponse when the result is empty. + */ + public void testSerializationNull() throws IOException { + List list = new ArrayList<>(); + GetQueryGroupResponse response = new GetQueryGroupResponse(list, RestStatus.OK); + assertEquals(response.getQueryGroups(), list); + + BytesStreamOutput out = new BytesStreamOutput(); + response.writeTo(out); + StreamInput streamInput = out.bytes().streamInput(); + + GetQueryGroupResponse otherResponse = new GetQueryGroupResponse(streamInput); + assertEquals(response.getRestStatus(), otherResponse.getRestStatus()); + assertEquals(0, otherResponse.getQueryGroups().size()); + } + + /** + * Test case to verify the toXContent of GetQueryGroupResponse. + */ + public void testToXContentGetSingleQueryGroup() throws IOException { + List queryGroupList = new ArrayList<>(); + queryGroupList.add(QueryGroupTestUtils.queryGroupOne); + XContentBuilder builder = JsonXContent.contentBuilder().prettyPrint(); + GetQueryGroupResponse response = new GetQueryGroupResponse(queryGroupList, RestStatus.OK); + String actual = response.toXContent(builder, mock(ToXContent.Params.class)).toString(); + String expected = "{\n" + + " \"query_groups\" : [\n" + + " {\n" + + " \"_id\" : \"AgfUO5Ja9yfsYlONlYi3TQ==\",\n" + + " \"name\" : \"query_group_one\",\n" + + " \"resiliency_mode\" : \"monitor\",\n" + + " \"updated_at\" : 4513232413,\n" + + " \"resource_limits\" : {\n" + + " \"memory\" : 0.3\n" + + " }\n" + + " }\n" + + " ]\n" + + "}"; + assertEquals(expected, actual); + } + + /** + * Test case to verify the toXContent of GetQueryGroupResponse when the result contains multiple QueryGroups. + */ + public void testToXContentGetMultipleQueryGroup() throws IOException { + List queryGroupList = new ArrayList<>(); + queryGroupList.add(QueryGroupTestUtils.queryGroupOne); + queryGroupList.add(QueryGroupTestUtils.queryGroupTwo); + XContentBuilder builder = JsonXContent.contentBuilder().prettyPrint(); + GetQueryGroupResponse response = new GetQueryGroupResponse(queryGroupList, RestStatus.OK); + String actual = response.toXContent(builder, mock(ToXContent.Params.class)).toString(); + String expected = "{\n" + + " \"query_groups\" : [\n" + + " {\n" + + " \"_id\" : \"AgfUO5Ja9yfsYlONlYi3TQ==\",\n" + + " \"name\" : \"query_group_one\",\n" + + " \"resiliency_mode\" : \"monitor\",\n" + + " \"updated_at\" : 4513232413,\n" + + " \"resource_limits\" : {\n" + + " \"memory\" : 0.3\n" + + " }\n" + + " },\n" + + " {\n" + + " \"_id\" : \"G5iIqHy4g7eK1qIAAAAIH53=1\",\n" + + " \"name\" : \"query_group_two\",\n" + + " \"resiliency_mode\" : \"monitor\",\n" + + " \"updated_at\" : 4513232415,\n" + + " \"resource_limits\" : {\n" + + " \"memory\" : 0.6\n" + + " }\n" + + " }\n" + + " ]\n" + + "}"; + assertEquals(expected, actual); + } + + /** + * Test case to verify toXContent of GetQueryGroupResponse when the result contains zero QueryGroup. + */ + public void testToXContentGetZeroQueryGroup() throws IOException { + XContentBuilder builder = JsonXContent.contentBuilder().prettyPrint(); + GetQueryGroupResponse otherResponse = new GetQueryGroupResponse(new ArrayList<>(), RestStatus.OK); + String actual = otherResponse.toXContent(builder, mock(ToXContent.Params.class)).toString(); + String expected = "{\n" + " \"query_groups\" : [ ]\n" + "}"; + assertEquals(expected, actual); + } +} diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/TransportGetQueryGroupActionTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/TransportGetQueryGroupActionTests.java new file mode 100644 index 0000000000000..755b11a5f4b89 --- /dev/null +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/TransportGetQueryGroupActionTests.java @@ -0,0 +1,47 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.plugin.wlm.action; + +import org.opensearch.ResourceNotFoundException; +import org.opensearch.action.support.ActionFilters; +import org.opensearch.cluster.metadata.IndexNameExpressionResolver; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.core.action.ActionListener; +import org.opensearch.test.OpenSearchTestCase; +import org.opensearch.threadpool.ThreadPool; +import org.opensearch.transport.TransportService; + +import static org.opensearch.plugin.wlm.QueryGroupTestUtils.NAME_NONE_EXISTED; +import static org.opensearch.plugin.wlm.QueryGroupTestUtils.NAME_ONE; +import static org.opensearch.plugin.wlm.QueryGroupTestUtils.clusterState; +import static org.mockito.Mockito.mock; + +public class TransportGetQueryGroupActionTests extends OpenSearchTestCase { + + /** + * Test case for ClusterManagerOperation function + */ + @SuppressWarnings("unchecked") + public void testClusterManagerOperation() throws Exception { + GetQueryGroupRequest getQueryGroupRequest1 = new GetQueryGroupRequest(NAME_NONE_EXISTED); + GetQueryGroupRequest getQueryGroupRequest2 = new GetQueryGroupRequest(NAME_ONE); + TransportGetQueryGroupAction transportGetQueryGroupAction = new TransportGetQueryGroupAction( + mock(ClusterService.class), + mock(TransportService.class), + mock(ActionFilters.class), + mock(ThreadPool.class), + mock(IndexNameExpressionResolver.class) + ); + assertThrows( + ResourceNotFoundException.class, + () -> transportGetQueryGroupAction.clusterManagerOperation(getQueryGroupRequest1, clusterState(), mock(ActionListener.class)) + ); + transportGetQueryGroupAction.clusterManagerOperation(getQueryGroupRequest2, clusterState(), mock(ActionListener.class)); + } +} diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceServiceTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceServiceTests.java index 533c98b44685d..2aa3b9e168852 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceServiceTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceServiceTests.java @@ -29,6 +29,8 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.Set; +import java.util.stream.Collectors; import org.mockito.ArgumentCaptor; @@ -42,6 +44,7 @@ import static org.opensearch.plugin.wlm.QueryGroupTestUtils.assertEqualQueryGroups; import static org.opensearch.plugin.wlm.QueryGroupTestUtils.clusterSettings; import static org.opensearch.plugin.wlm.QueryGroupTestUtils.clusterSettingsSet; +import static org.opensearch.plugin.wlm.QueryGroupTestUtils.clusterState; import static org.opensearch.plugin.wlm.QueryGroupTestUtils.preparePersistenceServiceSetup; import static org.opensearch.plugin.wlm.QueryGroupTestUtils.queryGroupList; import static org.opensearch.plugin.wlm.QueryGroupTestUtils.queryGroupOne; @@ -244,4 +247,55 @@ public void testPersistInClusterStateMetadataFailure() { queryGroupPersistenceService.persistInClusterStateMetadata(queryGroupOne, listener); verify(listener).onFailure(any(RuntimeException.class)); } + + /** + * Tests getting a single QueryGroup + */ + public void testGetSingleQueryGroup() { + Collection groupsCollections = QueryGroupPersistenceService.getFromClusterStateMetadata(NAME_ONE, clusterState()); + List groups = new ArrayList<>(groupsCollections); + assertEquals(1, groups.size()); + QueryGroup queryGroup = groups.get(0); + List listOne = new ArrayList<>(); + List listTwo = new ArrayList<>(); + listOne.add(QueryGroupTestUtils.queryGroupOne); + listTwo.add(queryGroup); + QueryGroupTestUtils.assertEqualQueryGroups(listOne, listTwo); + } + + /** + * Tests getting all QueryGroups + */ + public void testGetAllQueryGroups() { + assertEquals(2, QueryGroupTestUtils.clusterState().metadata().queryGroups().size()); + Collection groupsCollections = QueryGroupPersistenceService.getFromClusterStateMetadata(null, clusterState()); + List res = new ArrayList<>(groupsCollections); + assertEquals(2, res.size()); + Set currentNAME = res.stream().map(QueryGroup::getName).collect(Collectors.toSet()); + assertTrue(currentNAME.contains(QueryGroupTestUtils.NAME_ONE)); + assertTrue(currentNAME.contains(QueryGroupTestUtils.NAME_TWO)); + QueryGroupTestUtils.assertEqualQueryGroups(QueryGroupTestUtils.queryGroupList(), res); + } + + /** + * Tests getting a QueryGroup with invalid name + */ + public void testGetNonExistedQueryGroups() { + Collection groupsCollections = QueryGroupPersistenceService.getFromClusterStateMetadata( + NAME_NONE_EXISTED, + clusterState() + ); + List groups = new ArrayList<>(groupsCollections); + assertEquals(0, groups.size()); + } + + /** + * Tests setting maxQueryGroupCount + */ + public void testMaxQueryGroupCount() { + assertThrows(IllegalArgumentException.class, () -> QueryGroupTestUtils.queryGroupPersistenceService().setMaxQueryGroupCount(-1)); + QueryGroupPersistenceService queryGroupPersistenceService = QueryGroupTestUtils.queryGroupPersistenceService(); + queryGroupPersistenceService.setMaxQueryGroupCount(50); + assertEquals(50, queryGroupPersistenceService.getMaxQueryGroupCount()); + } } diff --git a/plugins/workload-management/src/yamlRestTest/resources/rest-api-spec/api/get_query_group_context.json b/plugins/workload-management/src/yamlRestTest/resources/rest-api-spec/api/get_query_group_context.json new file mode 100644 index 0000000000000..e0d552be616b2 --- /dev/null +++ b/plugins/workload-management/src/yamlRestTest/resources/rest-api-spec/api/get_query_group_context.json @@ -0,0 +1,25 @@ +{ + "get_query_group_context": { + "stability": "experimental", + "url": { + "paths": [ + { + "path": "/_wlm/query_group", + "methods": ["GET"], + "parts": {} + }, + { + "path": "/_wlm/query_group/{name}", + "methods": ["GET"], + "parts": { + "name": { + "type": "string", + "description": "QueryGroup name" + } + } + } + ] + }, + "params":{} + } +} diff --git a/plugins/workload-management/src/yamlRestTest/resources/rest-api-spec/test/wlm/10_create_query_group.yml b/plugins/workload-management/src/yamlRestTest/resources/rest-api-spec/test/wlm/10_query_group.yml similarity index 77% rename from plugins/workload-management/src/yamlRestTest/resources/rest-api-spec/test/wlm/10_create_query_group.yml rename to plugins/workload-management/src/yamlRestTest/resources/rest-api-spec/test/wlm/10_query_group.yml index ae82a8146e9cd..a22dfa2f4477e 100644 --- a/plugins/workload-management/src/yamlRestTest/resources/rest-api-spec/test/wlm/10_create_query_group.yml +++ b/plugins/workload-management/src/yamlRestTest/resources/rest-api-spec/test/wlm/10_query_group.yml @@ -1,4 +1,4 @@ -"test create QueryGroup API": +"test CRUD Operations for QueryGroup API ": - skip: version: " - 2.16.99" reason: "QueryGroup WorkloadManagement feature was added in 2.17" @@ -20,6 +20,15 @@ - match: { resource_limits.cpu: 0.4 } - match: { resource_limits.memory: 0.2 } + - do: + get_query_group_context: + name: "analytics" + + - match: { query_groups.0.name: "analytics" } + - match: { query_groups.0.resiliency_mode: "monitor" } + - match: { query_groups.0.resource_limits.cpu: 0.4 } + - match: { query_groups.0.resource_limits.memory: 0.2 } + - do: catch: /illegal_argument_exception/ create_query_group_context: @@ -88,3 +97,12 @@ - match: { resiliency_mode: "monitor" } - match: { resource_limits.cpu: 0.35 } - match: { resource_limits.memory: 0.25 } + + - do: + get_query_group_context: + name: "analytics2" + + - match: { query_groups.0.name: "analytics2" } + - match: { query_groups.0.resiliency_mode: "monitor" } + - match: { query_groups.0.resource_limits.cpu: 0.35 } + - match: { query_groups.0.resource_limits.memory: 0.25 } diff --git a/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java b/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java index 6ab11b1d6f150..9b5c6bc2369a6 100644 --- a/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java +++ b/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java @@ -62,10 +62,7 @@ public QueryGroup(String name, String _id, ResiliencyMode resiliencyMode, Map MAX_CHARS_ALLOWED_IN_NAME || name.isEmpty()) { - throw new IllegalArgumentException("QueryGroup.name shouldn't be empty or more than 50 chars long"); - } + validateName(name); if (resourceLimits.isEmpty()) { throw new IllegalArgumentException("QueryGroup.resourceLimits should at least have 1 resource limit"); @@ -111,6 +108,12 @@ public void writeTo(StreamOutput out) throws IOException { out.writeLong(updatedAtInMillis); } + public static void validateName(String name) { + if (name == null || name.isEmpty() || name.length() > MAX_CHARS_ALLOWED_IN_NAME) { + throw new IllegalArgumentException("QueryGroup.name shouldn't be null, empty or more than 50 chars long"); + } + } + private void validateResourceLimits(Map resourceLimits) { for (Map.Entry resource : resourceLimits.entrySet()) { Double threshold = resource.getValue(); From 9661e8db62e8897cb53b0bc1860ead43b28dfa21 Mon Sep 17 00:00:00 2001 From: Matthew Ridehalgh <9155962+mridehalgh@users.noreply.github.com> Date: Fri, 16 Aug 2024 20:11:13 +0100 Subject: [PATCH 40/40] fix: MinHash token filter parameters not working (#15233) * fix: minhash configuration Signed-off-by: Matt Ridehalgh * style: linting Signed-off-by: Matt Ridehalgh * chore: update CHANGELOG.md Signed-off-by: Matt Ridehalgh --------- Signed-off-by: Matt Ridehalgh --- CHANGELOG.md | 3 +- .../common/MinHashTokenFilterFactory.java | 4 +- .../common/MinHashFilterFactoryTests.java | 66 +++++++++++++++---- 3 files changed, 58 insertions(+), 15 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9fe700d7190fa..5103d978b361a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -34,7 +34,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), - Bump `com.azure:azure-core` from 1.49.1 to 1.51.0 ([#15111](https://github.com/opensearch-project/OpenSearch/pull/15111)) - Bump `org.xerial.snappy:snappy-java` from 1.1.10.5 to 1.1.10.6 ([#15207](https://github.com/opensearch-project/OpenSearch/pull/15207)) - Bump `com.azure:azure-xml` from 1.0.0 to 1.1.0 ([#15206](https://github.com/opensearch-project/OpenSearch/pull/15206)) -- Bump `reactor` from 3.5.19 to 3.5.20 ([#15262](https://github.com/opensearch-project/OpenSearch/pull/15262)) +- Bump `reactor` from 3.5.19 to 3.5.20 ([#15262](https://github.com/opensearch-project/OpenSearch/pull/15262)) - Bump `reactor-netty` from 1.1.21 to 1.1.22 ([#15262](https://github.com/opensearch-project/OpenSearch/pull/15262)) ### Changed @@ -53,6 +53,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), - Fix array_index_out_of_bounds_exception when indexing documents with field name containing only dot ([#15126](https://github.com/opensearch-project/OpenSearch/pull/15126)) - Fixed array field name omission in flat_object function for nested JSON ([#13620](https://github.com/opensearch-project/OpenSearch/pull/13620)) - Fix range aggregation optimization ignoring top level queries ([#15194](https://github.com/opensearch-project/OpenSearch/pull/15194)) +- Fix incorrect parameter names in MinHash token filter configuration handling ([#15233](https://github.com/opensearch-project/OpenSearch/pull/15233)) ### Security diff --git a/modules/analysis-common/src/main/java/org/opensearch/analysis/common/MinHashTokenFilterFactory.java b/modules/analysis-common/src/main/java/org/opensearch/analysis/common/MinHashTokenFilterFactory.java index e76354ae3a765..40655b84794d5 100644 --- a/modules/analysis-common/src/main/java/org/opensearch/analysis/common/MinHashTokenFilterFactory.java +++ b/modules/analysis-common/src/main/java/org/opensearch/analysis/common/MinHashTokenFilterFactory.java @@ -65,10 +65,10 @@ private Map convertSettings(Settings settings) { if (settings.hasValue("hash_count")) { settingMap.put("hashCount", settings.get("hash_count")); } - if (settings.hasValue("bucketCount")) { + if (settings.hasValue("bucket_count")) { settingMap.put("bucketCount", settings.get("bucket_count")); } - if (settings.hasValue("hashSetSize")) { + if (settings.hasValue("hash_set_size")) { settingMap.put("hashSetSize", settings.get("hash_set_size")); } if (settings.hasValue("with_rotation")) { diff --git a/modules/analysis-common/src/test/java/org/opensearch/analysis/common/MinHashFilterFactoryTests.java b/modules/analysis-common/src/test/java/org/opensearch/analysis/common/MinHashFilterFactoryTests.java index e86a939dc857b..b9f09033f49f3 100644 --- a/modules/analysis-common/src/test/java/org/opensearch/analysis/common/MinHashFilterFactoryTests.java +++ b/modules/analysis-common/src/test/java/org/opensearch/analysis/common/MinHashFilterFactoryTests.java @@ -50,14 +50,10 @@ public void testDefault() throws IOException { int default_bucket_size = 512; int default_hash_set_size = 1; Settings settings = Settings.builder().put(Environment.PATH_HOME_SETTING.getKey(), createTempDir().toString()).build(); - OpenSearchTestCase.TestAnalysis analysis = AnalysisTestsHelper.createTestAnalysisFromSettings( - settings, - new CommonAnalysisModulePlugin() - ); + OpenSearchTestCase.TestAnalysis analysis = getTestAnalysisFromSettings(settings); TokenFilterFactory tokenFilter = analysis.tokenFilter.get("min_hash"); String source = "the quick brown fox"; - Tokenizer tokenizer = new WhitespaceTokenizer(); - tokenizer.setReader(new StringReader(source)); + Tokenizer tokenizer = getTokenizer(source); // with_rotation is true by default, and hash_set_size is 1, so even though the source doesn't // have enough tokens to fill all the buckets, we still expect 512 tokens. @@ -73,17 +69,63 @@ public void testSettings() throws IOException { .put("index.analysis.filter.test_min_hash.with_rotation", false) .put(Environment.PATH_HOME_SETTING.getKey(), createTempDir().toString()) .build(); - OpenSearchTestCase.TestAnalysis analysis = AnalysisTestsHelper.createTestAnalysisFromSettings( - settings, - new CommonAnalysisModulePlugin() - ); + OpenSearchTestCase.TestAnalysis analysis = getTestAnalysisFromSettings(settings); TokenFilterFactory tokenFilter = analysis.tokenFilter.get("test_min_hash"); String source = "sushi"; - Tokenizer tokenizer = new WhitespaceTokenizer(); - tokenizer.setReader(new StringReader(source)); + Tokenizer tokenizer = getTokenizer(source); // despite the fact that bucket_count is 2 and hash_set_size is 1, // because with_rotation is false, we only expect 1 token here. assertStreamHasNumberOfTokens(tokenFilter.create(tokenizer), 1); } + + public void testBucketCountSetting() throws IOException { + // Correct case with "bucket_count" + Settings settingsWithBucketCount = Settings.builder() + .put("index.analysis.filter.test_min_hash.type", "min_hash") + .put("index.analysis.filter.test_min_hash.hash_count", "1") + .put("index.analysis.filter.test_min_hash.bucket_count", "3") + .put("index.analysis.filter.test_min_hash.hash_set_size", "1") + .put("index.analysis.filter.test_min_hash.with_rotation", false) + .put(Environment.PATH_HOME_SETTING.getKey(), createTempDir().toString()) + .build(); + + OpenSearchTestCase.TestAnalysis analysisWithBucketCount = getTestAnalysisFromSettings(settingsWithBucketCount); + + TokenFilterFactory tokenFilterWithBucketCount = analysisWithBucketCount.tokenFilter.get("test_min_hash"); + String sourceWithBucketCount = "salmon avocado roll uramaki"; + Tokenizer tokenizerWithBucketCount = getTokenizer(sourceWithBucketCount); + // Expect 3 tokens due to bucket_count being set to 3 + assertStreamHasNumberOfTokens(tokenFilterWithBucketCount.create(tokenizerWithBucketCount), 3); + } + + public void testHashSetSizeSetting() throws IOException { + // Correct case with "hash_set_size" + Settings settingsWithHashSetSize = Settings.builder() + .put("index.analysis.filter.test_min_hash.type", "min_hash") + .put("index.analysis.filter.test_min_hash.hash_count", "1") + .put("index.analysis.filter.test_min_hash.bucket_count", "1") + .put("index.analysis.filter.test_min_hash.hash_set_size", "2") + .put("index.analysis.filter.test_min_hash.with_rotation", false) + .put(Environment.PATH_HOME_SETTING.getKey(), createTempDir().toString()) + .build(); + + OpenSearchTestCase.TestAnalysis analysisWithHashSetSize = getTestAnalysisFromSettings(settingsWithHashSetSize); + + TokenFilterFactory tokenFilterWithHashSetSize = analysisWithHashSetSize.tokenFilter.get("test_min_hash"); + String sourceWithHashSetSize = "salmon avocado roll uramaki"; + Tokenizer tokenizerWithHashSetSize = getTokenizer(sourceWithHashSetSize); + // Expect 2 tokens due to hash_set_size being set to 2 and bucket_count being 1 + assertStreamHasNumberOfTokens(tokenFilterWithHashSetSize.create(tokenizerWithHashSetSize), 2); + } + + private static OpenSearchTestCase.TestAnalysis getTestAnalysisFromSettings(Settings settingsWithBucketCount) throws IOException { + return AnalysisTestsHelper.createTestAnalysisFromSettings(settingsWithBucketCount, new CommonAnalysisModulePlugin()); + } + + private static Tokenizer getTokenizer(String sourceWithBucketCount) { + Tokenizer tokenizerWithBucketCount = new WhitespaceTokenizer(); + tokenizerWithBucketCount.setReader(new StringReader(sourceWithBucketCount)); + return tokenizerWithBucketCount; + } }