Skip to content

Commit 0287c7a

Browse files
dzane17dzane17
committed
WLM group custom search settings - groundwork and timeout
Signed-off-by: David Zane <davizane@amazon.com>
1 parent a701344 commit 0287c7a

File tree

15 files changed

+647
-56
lines changed

15 files changed

+647
-56
lines changed

‎CHANGELOG.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
3838
- Introduced strategy planner interfaces for indexing and deletion ([#20585](https://github.com/opensearch-project/OpenSearch/pull/20585))
3939
- Implement FieldMappingIngestionMessageMapper for pull-based ingestion ([#20729](https://github.com/opensearch-project/OpenSearch/pull/20729))
4040
- Added support of WarmerRefreshListener in NRTReplicationEngine to trigger warmer after replication on replica shards ([#20650](https://github.com/opensearch-project/OpenSearch/pull/20650))
41+
- WLM group custom search settings - groundwork and timeout ([#20536](https://github.com/opensearch-project/OpenSearch/issues/20536))
4142

4243
### Changed
4344
- Make telemetry `Tags` immutable ([#20788](https://github.com/opensearch-project/OpenSearch/pull/20788))

‎plugins/workload-management/src/javaRestTest/java/org/opensearch/rest/WorkloadManagementRestIT.java‎

Lines changed: 62 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -27,16 +27,16 @@ public void enableWlmMode() throws Exception {
2727

2828
public void testCreate() throws Exception {
2929
Response response = performOperation("PUT", "_wlm/workload_group", getCreateJson("analytics", "enforced", 0.4, 0.2));
30-
assertEquals(response.getStatusLine().getStatusCode(), 200);
30+
assertEquals(200, response.getStatusLine().getStatusCode());
3131
performOperation("DELETE", "_wlm/workload_group/analytics", null);
3232
}
3333

3434
public void testMultipleCreate() throws Exception {
3535
Response response = performOperation("PUT", "_wlm/workload_group", getCreateJson("analytics2", "enforced", 0.4, 0.2));
36-
assertEquals(response.getStatusLine().getStatusCode(), 200);
36+
assertEquals(200, response.getStatusLine().getStatusCode());
3737

3838
Response response2 = performOperation("PUT", "_wlm/workload_group", getCreateJson("users", "soft", 0.2, 0.1));
39-
assertEquals(response2.getStatusLine().getStatusCode(), 200);
39+
assertEquals(200, response2.getStatusLine().getStatusCode());
4040

4141
assertThrows(
4242
ResponseException.class,
@@ -48,10 +48,10 @@ public void testMultipleCreate() throws Exception {
4848

4949
public void testGet() throws Exception {
5050
Response response = performOperation("PUT", "_wlm/workload_group", getCreateJson("analytics3", "enforced", 0.4, 0.2));
51-
assertEquals(response.getStatusLine().getStatusCode(), 200);
51+
assertEquals(200, response.getStatusLine().getStatusCode());
5252

5353
Response response2 = performOperation("GET", "_wlm/workload_group/analytics3", null);
54-
assertEquals(response2.getStatusLine().getStatusCode(), 200);
54+
assertEquals(200, response2.getStatusLine().getStatusCode());
5555
String responseBody2 = EntityUtils.toString(response2.getEntity());
5656
assertTrue(responseBody2.contains("\"name\":\"analytics3\""));
5757
assertTrue(responseBody2.contains("\"resiliency_mode\":\"enforced\""));
@@ -64,33 +64,34 @@ public void testGet() throws Exception {
6464

6565
public void testDelete() throws Exception {
6666
Response response = performOperation("PUT", "_wlm/workload_group", getCreateJson("analytics4", "enforced", 0.4, 0.2));
67-
assertEquals(response.getStatusLine().getStatusCode(), 200);
67+
assertEquals(200, response.getStatusLine().getStatusCode());
6868

6969
Response response2 = performOperation("DELETE", "_wlm/workload_group/analytics4", null);
70-
assertEquals(response2.getStatusLine().getStatusCode(), 200);
70+
assertEquals(200, response2.getStatusLine().getStatusCode());
7171
assertTrue(EntityUtils.toString(response2.getEntity()).contains("\"acknowledged\":true"));
7272

7373
assertThrows(ResponseException.class, () -> performOperation("DELETE", "_wlm/workload_group/analytics99", null));
7474
}
7575

7676
public void testUpdate() throws Exception {
7777
Response response = performOperation("PUT", "_wlm/workload_group", getCreateJson("analytics5", "enforced", 0.4, 0.2));
78-
assertEquals(response.getStatusLine().getStatusCode(), 200);
78+
assertEquals(200, response.getStatusLine().getStatusCode());
7979

8080
Response response2 = performOperation("PUT", "_wlm/workload_group/analytics5", getUpdateJson("monitor", 0.41, 0.21));
81-
assertEquals(response2.getStatusLine().getStatusCode(), 200);
81+
assertEquals(200, response2.getStatusLine().getStatusCode());
8282
String responseBody2 = EntityUtils.toString(response2.getEntity());
8383
assertTrue(responseBody2.contains("\"name\":\"analytics5\""));
8484
assertTrue(responseBody2.contains("\"resiliency_mode\":\"monitor\""));
8585
assertTrue(responseBody2.contains("\"cpu\":0.41"));
8686
assertTrue(responseBody2.contains("\"memory\":0.21"));
8787

88-
String json = "{\n"
89-
+ " \"resource_limits\": {\n"
90-
+ " \"cpu\" : 1.1,\n"
91-
+ " \"memory\" : -0.1\n"
92-
+ " }\n"
93-
+ "}'";
88+
String json = """
89+
{
90+
"resource_limits": {
91+
"cpu" : 1.1,
92+
"memory" : -0.1
93+
}
94+
}""";
9495
assertThrows(ResponseException.class, () -> performOperation("PUT", "_wlm/workload_group/analytics5", json));
9596
assertThrows(
9697
ResponseException.class,
@@ -101,13 +102,13 @@ public void testUpdate() throws Exception {
101102

102103
public void testCRUD() throws Exception {
103104
Response response = performOperation("PUT", "_wlm/workload_group", getCreateJson("analytics6", "enforced", 0.4, 0.2));
104-
assertEquals(response.getStatusLine().getStatusCode(), 200);
105+
assertEquals(200, response.getStatusLine().getStatusCode());
105106

106107
Response response1 = performOperation("PUT", "_wlm/workload_group/analytics6", getUpdateJson("monitor", 0.41, 0.21));
107-
assertEquals(response1.getStatusLine().getStatusCode(), 200);
108+
assertEquals(200, response1.getStatusLine().getStatusCode());
108109

109110
Response response2 = performOperation("GET", "_wlm/workload_group/analytics6", null);
110-
assertEquals(response2.getStatusLine().getStatusCode(), 200);
111+
assertEquals(200, response2.getStatusLine().getStatusCode());
111112
String responseBody2 = EntityUtils.toString(response2.getEntity());
112113
assertTrue(responseBody2.contains("\"name\":\"analytics6\""));
113114
assertTrue(responseBody2.contains("\"resiliency_mode\":\"monitor\""));
@@ -120,15 +121,15 @@ public void testCRUD() throws Exception {
120121
);
121122

122123
Response response4 = performOperation("PUT", "_wlm/workload_group", getCreateJson("users3", "monitor", 0.59, 0.79));
123-
assertEquals(response4.getStatusLine().getStatusCode(), 200);
124+
assertEquals(200, response4.getStatusLine().getStatusCode());
124125

125126
Response response5 = performOperation("DELETE", "_wlm/workload_group/analytics6", null);
126-
assertEquals(response5.getStatusLine().getStatusCode(), 200);
127+
assertEquals(200, response5.getStatusLine().getStatusCode());
127128
String responseBody5 = EntityUtils.toString(response5.getEntity());
128129
assertTrue(responseBody5.contains("\"acknowledged\":true"));
129130

130131
Response response6 = performOperation("GET", "_wlm/workload_group", null);
131-
assertEquals(response6.getStatusLine().getStatusCode(), 200);
132+
assertEquals(200, response6.getStatusLine().getStatusCode());
132133
String responseBody6 = EntityUtils.toString(response6.getEntity());
133134
assertTrue(responseBody6.contains("\"workload_groups\""));
134135
assertTrue(responseBody6.contains("\"users3\""));
@@ -146,6 +147,44 @@ public void testOperationWhenWlmDisabled() throws Exception {
146147
assertOK(performOperation("GET", "_wlm/workload_group/", null));
147148
}
148149

150+
public void testSearchSettings() throws Exception {
151+
// Create with search_settings
152+
String createJson = """
153+
{
154+
"name": "search_test",
155+
"resiliency_mode": "enforced",
156+
"resource_limits": {"cpu": 0.3, "memory": 0.3},
157+
"search_settings": {
158+
"timeout": "30s"
159+
}
160+
}""";
161+
Response response = performOperation("PUT", "_wlm/workload_group", createJson);
162+
assertEquals(200, response.getStatusLine().getStatusCode());
163+
164+
// Verify search_settings in GET response
165+
Response getResponse = performOperation("GET", "_wlm/workload_group/search_test", null);
166+
String responseBody = EntityUtils.toString(getResponse.getEntity());
167+
assertTrue(responseBody.contains("\"search_settings\""));
168+
assertTrue(responseBody.contains("\"timeout\":\"30s\""));
169+
170+
// Update search_settings
171+
String updateJson = """
172+
{
173+
"search_settings": {
174+
"timeout": "1m"
175+
}
176+
}""";
177+
Response updateResponse = performOperation("PUT", "_wlm/workload_group/search_test", updateJson);
178+
assertEquals(200, updateResponse.getStatusLine().getStatusCode());
179+
180+
// Verify updated search_settings
181+
Response getResponse2 = performOperation("GET", "_wlm/workload_group/search_test", null);
182+
String responseBody2 = EntityUtils.toString(getResponse2.getEntity());
183+
assertTrue(responseBody2.contains("\"timeout\":\"1m\""));
184+
185+
performOperation("DELETE", "_wlm/workload_group/search_test", null);
186+
}
187+
149188
static String getCreateJson(String name, String resiliencyMode, double cpu, double memory) {
150189
return "{\n"
151190
+ " \"name\": \""
@@ -161,7 +200,8 @@ static String getCreateJson(String name, String resiliencyMode, double cpu, doub
161200
+ " \"memory\" : "
162201
+ memory
163202
+ "\n"
164-
+ " }\n"
203+
+ " },\n"
204+
+ " \"search_settings\": {}\n"
165205
+ "}";
166206
}
167207

‎plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/WorkloadManagementTestUtils.java‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,11 +43,15 @@
4343
public class WorkloadManagementTestUtils {
4444
public static final String NAME_ONE = "workload_group_one";
4545
public static final String NAME_TWO = "workload_group_two";
46+
public static final String NAME_THREE = "workload_group_three";
4647
public static final String _ID_ONE = "AgfUO5Ja9yfsYlONlYi3TQ==";
4748
public static final String _ID_TWO = "G5iIqHy4g7eK1qIAAAAIH53=1";
49+
public static final String _ID_THREE = "H6jVP6Kb0zgtZmPOmZj4UQ==";
4850
public static final String NAME_NONE_EXISTED = "workload_group_none_existed";
4951
public static final long TIMESTAMP_ONE = 4513232413L;
5052
public static final long TIMESTAMP_TWO = 4513232415L;
53+
public static final long TIMESTAMP_THREE = 4513232417L;
54+
public static final Map<String, String> TEST_SEARCH_SETTINGS = Map.of("timeout", "30s");
5155
public static final WorkloadGroup workloadGroupOne = builder().name(NAME_ONE)
5256
._id(_ID_ONE)
5357
.mutableWorkloadGroupFragment(
@@ -64,6 +68,18 @@ public class WorkloadManagementTestUtils {
6468
.updatedAt(TIMESTAMP_TWO)
6569
.build();
6670

71+
public static final WorkloadGroup workloadGroupWithSearchSettings = builder().name(NAME_THREE)
72+
._id(_ID_THREE)
73+
.mutableWorkloadGroupFragment(
74+
new MutableWorkloadGroupFragment(
75+
MutableWorkloadGroupFragment.ResiliencyMode.ENFORCED,
76+
Map.of(ResourceType.MEMORY, 0.5),
77+
TEST_SEARCH_SETTINGS
78+
)
79+
)
80+
.updatedAt(TIMESTAMP_THREE)
81+
.build();
82+
6783
public static List<WorkloadGroup> workloadGroupList() {
6884
List<WorkloadGroup> list = new ArrayList<>();
6985
list.add(workloadGroupOne);

‎plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateWorkloadGroupResponseTests.java‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,8 +59,34 @@ public void testToXContentCreateWorkloadGroup() throws IOException {
5959
+ " \"resource_limits\" : {\n"
6060
+ " \"memory\" : 0.3\n"
6161
+ " },\n"
62+
+ " \"search_settings\" : { },\n"
6263
+ " \"updated_at\" : 4513232413\n"
6364
+ "}";
6465
assertEquals(expected, actual);
6566
}
67+
68+
/**
69+
* Test case to validate the toXContent method of CreateWorkloadGroupResponse with search settings.
70+
*/
71+
public void testToXContentCreateWorkloadGroupWithSearchSettings() throws IOException {
72+
XContentBuilder builder = JsonXContent.contentBuilder().prettyPrint();
73+
CreateWorkloadGroupResponse response = new CreateWorkloadGroupResponse(
74+
WorkloadManagementTestUtils.workloadGroupWithSearchSettings,
75+
RestStatus.OK
76+
);
77+
String actual = response.toXContent(builder, mock(ToXContent.Params.class)).toString();
78+
String expected = "{\n"
79+
+ " \"_id\" : \"H6jVP6Kb0zgtZmPOmZj4UQ==\",\n"
80+
+ " \"name\" : \"workload_group_three\",\n"
81+
+ " \"resiliency_mode\" : \"enforced\",\n"
82+
+ " \"resource_limits\" : {\n"
83+
+ " \"memory\" : 0.5\n"
84+
+ " },\n"
85+
+ " \"search_settings\" : {\n"
86+
+ " \"timeout\" : \"30s\"\n"
87+
+ " },\n"
88+
+ " \"updated_at\" : 4513232417\n"
89+
+ "}";
90+
assertEquals(expected, actual);
91+
}
6692
}

‎plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/GetWorkloadGroupResponseTests.java‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -97,6 +97,7 @@ public void testToXContentGetSingleWorkloadGroup() throws IOException {
9797
"resource_limits" : {
9898
"memory" : 0.3
9999
},
100+
"search_settings" : { },
100101
"updated_at" : 4513232413
101102
}
102103
]
@@ -124,6 +125,7 @@ public void testToXContentGetMultipleWorkloadGroup() throws IOException {
124125
"resource_limits" : {
125126
"memory" : 0.3
126127
},
128+
"search_settings" : { },
127129
"updated_at" : 4513232413
128130
},
129131
{
@@ -133,6 +135,7 @@ public void testToXContentGetMultipleWorkloadGroup() throws IOException {
133135
"resource_limits" : {
134136
"memory" : 0.6
135137
},
138+
"search_settings" : { },
136139
"updated_at" : 4513232415
137140
}
138141
]
@@ -153,4 +156,33 @@ public void testToXContentGetZeroWorkloadGroup() throws IOException {
153156
}""";
154157
assertEquals(expected, actual);
155158
}
159+
160+
/**
161+
* Test case to verify toXContent of GetWorkloadGroupResponse with search settings.
162+
*/
163+
public void testToXContentGetWorkloadGroupWithSearchSettings() throws IOException {
164+
List<WorkloadGroup> workloadGroupList = new ArrayList<>();
165+
workloadGroupList.add(WorkloadManagementTestUtils.workloadGroupWithSearchSettings);
166+
XContentBuilder builder = JsonXContent.contentBuilder().prettyPrint();
167+
GetWorkloadGroupResponse response = new GetWorkloadGroupResponse(workloadGroupList, RestStatus.OK);
168+
String actual = response.toXContent(builder, mock(ToXContent.Params.class)).toString();
169+
String expected = """
170+
{
171+
"workload_groups" : [
172+
{
173+
"_id" : "H6jVP6Kb0zgtZmPOmZj4UQ==",
174+
"name" : "workload_group_three",
175+
"resiliency_mode" : "enforced",
176+
"resource_limits" : {
177+
"memory" : 0.5
178+
},
179+
"search_settings" : {
180+
"timeout" : "30s"
181+
},
182+
"updated_at" : 4513232417
183+
}
184+
]
185+
}""";
186+
assertEquals(expected, actual);
187+
}
156188
}

‎plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/UpdateWorkloadGroupResponseTests.java‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -60,8 +60,34 @@ public void testToXContentUpdateSingleWorkloadGroup() throws IOException {
6060
+ " \"resource_limits\" : {\n"
6161
+ " \"memory\" : 0.3\n"
6262
+ " },\n"
63+
+ " \"search_settings\" : { },\n"
6364
+ " \"updated_at\" : 4513232413\n"
6465
+ "}";
6566
assertEquals(expected, actual);
6667
}
68+
69+
/**
70+
* Test case to verify the toXContent method of UpdateWorkloadGroupResponse with search settings.
71+
*/
72+
public void testToXContentUpdateWorkloadGroupWithSearchSettings() throws IOException {
73+
XContentBuilder builder = JsonXContent.contentBuilder().prettyPrint();
74+
UpdateWorkloadGroupResponse response = new UpdateWorkloadGroupResponse(
75+
WorkloadManagementTestUtils.workloadGroupWithSearchSettings,
76+
RestStatus.OK
77+
);
78+
String actual = response.toXContent(builder, mock(ToXContent.Params.class)).toString();
79+
String expected = "{\n"
80+
+ " \"_id\" : \"H6jVP6Kb0zgtZmPOmZj4UQ==\",\n"
81+
+ " \"name\" : \"workload_group_three\",\n"
82+
+ " \"resiliency_mode\" : \"enforced\",\n"
83+
+ " \"resource_limits\" : {\n"
84+
+ " \"memory\" : 0.5\n"
85+
+ " },\n"
86+
+ " \"search_settings\" : {\n"
87+
+ " \"timeout\" : \"30s\"\n"
88+
+ " },\n"
89+
+ " \"updated_at\" : 4513232417\n"
90+
+ "}";
91+
assertEquals(expected, actual);
92+
}
6793
}

‎server/src/main/java/org/opensearch/cluster/metadata/WorkloadGroup.java‎

Lines changed: 27 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,15 @@ public WorkloadGroup(String name, String _id, MutableWorkloadGroupFragment mutab
7373
throw new IllegalArgumentException("WorkloadGroup.updatedAtInMillis is not a valid epoch");
7474
}
7575

76+
// Normalize null searchSettings to empty map for storage
77+
if (mutableWorkloadGroupFragment.getSearchSettings() == null) {
78+
mutableWorkloadGroupFragment = new MutableWorkloadGroupFragment(
79+
mutableWorkloadGroupFragment.getResiliencyMode(),
80+
mutableWorkloadGroupFragment.getResourceLimits(),
81+
new HashMap<>()
82+
);
83+
}
84+
7685
this.name = name;
7786
this._id = _id;
7887
this.mutableWorkloadGroupFragment = mutableWorkloadGroupFragment;
@@ -104,10 +113,23 @@ public static WorkloadGroup updateExistingWorkloadGroup(
104113
}
105114
final ResiliencyMode mode = Optional.ofNullable(mutableWorkloadGroupFragment.getResiliencyMode())
106115
.orElse(existingGroup.getResiliencyMode());
116+
// Handle search_settings update:
117+
// null = not specified (keep existing)
118+
// empty map = explicitly clear (set to empty)
119+
// non-empty map = replace with new values
120+
final Map<String, String> mutableFragmentSearchSettings = mutableWorkloadGroupFragment.getSearchSettings();
121+
final Map<String, String> updatedSearchSettings;
122+
if (mutableFragmentSearchSettings == null) {
123+
// Not specified - keep existing
124+
updatedSearchSettings = new HashMap<>(existingGroup.getSearchSettings());
125+
} else {
126+
// Specified (empty or non-empty) - use the new value
127+
updatedSearchSettings = new HashMap<>(mutableFragmentSearchSettings);
128+
}
107129
return new WorkloadGroup(
108130
existingGroup.getName(),
109131
existingGroup.get_id(),
110-
new MutableWorkloadGroupFragment(mode, updatedResourceLimits),
132+
new MutableWorkloadGroupFragment(mode, updatedResourceLimits, updatedSearchSettings),
111133
Instant.now().getMillis()
112134
);
113135
}
@@ -179,6 +201,10 @@ public Map<ResourceType, Double> getResourceLimits() {
179201
return getMutableWorkloadGroupFragment().getResourceLimits();
180202
}
181203

204+
public Map<String, String> getSearchSettings() {
205+
return getMutableWorkloadGroupFragment().getSearchSettings();
206+
}
207+
182208
public String get_id() {
183209
return _id;
184210
}

0 commit comments

Comments
 (0)