Skip to content

Commit f03de82

Browse files
nik9000nik9000
committed
REST Client: NodeSelector for node attributes (#31296)
Add a `NodeSelector` so that users can filter the nodes that receive requests based on node attributes. I believe we'll need this to backport #30523 and we want it anyway. I also added a bash script to help with rebuilding the sniffer parsing test documents.
1 parent 3c0a375 commit f03de82

File tree

20 files changed

+1110
-476
lines changed

20 files changed

+1110
-476
lines changed

‎client/rest/src/main/java/org/elasticsearch/client/HasAttributeNodeSelector.java‎

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
/*
2+
* Licensed to Elasticsearch under one or more contributor
3+
* license agreements. See the NOTICE file distributed with
4+
* this work for additional information regarding copyright
5+
* ownership. Elasticsearch licenses this file to you under
6+
* the Apache License, Version 2.0 (the "License"); you may
7+
* not use this file except in compliance with the License.
8+
* You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
package org.elasticsearch.client;
21+
22+
import java.util.Iterator;
23+
import java.util.List;
24+
import java.util.Map;
25+
26+
/**
27+
* A {@link NodeSelector} that selects nodes that have a particular value
28+
* for an attribute.
29+
*/
30+
public final class HasAttributeNodeSelector implements NodeSelector {
31+
private final String key;
32+
private final String value;
33+
34+
public HasAttributeNodeSelector(String key, String value) {
35+
this.key = key;
36+
this.value = value;
37+
}
38+
39+
@Override
40+
public void select(Iterable<Node> nodes) {
41+
Iterator<Node> itr = nodes.iterator();
42+
while (itr.hasNext()) {
43+
Map<String, List<String>> allAttributes = itr.next().getAttributes();
44+
if (allAttributes == null) continue;
45+
List<String> values = allAttributes.get(key);
46+
if (values == null || false == values.contains(value)) {
47+
itr.remove();
48+
}
49+
}
50+
}
51+
52+
@Override
53+
public String toString() {
54+
return key + "=" + value;
55+
}
56+
}

‎client/rest/src/main/java/org/elasticsearch/client/Node.java‎

Lines changed: 23 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,8 @@
1919

2020
package org.elasticsearch.client;
2121

22+
import java.util.List;
23+
import java.util.Map;
2224
import java.util.Objects;
2325
import java.util.Set;
2426

@@ -52,13 +54,18 @@ public class Node {
5254
* if we don't know what roles the node has.
5355
*/
5456
private final Roles roles;
57+
/**
58+
* Attributes declared on the node.
59+
*/
60+
private final Map<String, List<String>> attributes;
5561

5662
/**
5763
* Create a {@linkplain Node} with metadata. All parameters except
5864
* {@code host} are nullable and implementations of {@link NodeSelector}
5965
* need to decide what to do in their absence.
6066
*/
61-
public Node(HttpHost host, Set<HttpHost> boundHosts, String name, String version, Roles roles) {
67+
public Node(HttpHost host, Set<HttpHost> boundHosts, String name, String version,
68+
Roles roles, Map<String, List<String>> attributes) {
6269
if (host == null) {
6370
throw new IllegalArgumentException("host cannot be null");
6471
}
@@ -67,13 +74,14 @@ public Node(HttpHost host, Set<HttpHost> boundHosts, String name, String version
6774
this.name = name;
6875
this.version = version;
6976
this.roles = roles;
77+
this.attributes = attributes;
7078
}
7179

7280
/**
7381
* Create a {@linkplain Node} without any metadata.
7482
*/
7583
public Node(HttpHost host) {
76-
this(host, null, null, null, null);
84+
this(host, null, null, null, null, null);
7785
}
7886

7987
/**
@@ -115,6 +123,13 @@ public Roles getRoles() {
115123
return roles;
116124
}
117125

126+
/**
127+
* Attributes declared on the node.
128+
*/
129+
public Map<String, List<String>> getAttributes() {
130+
return attributes;
131+
}
132+
118133
@Override
119134
public String toString() {
120135
StringBuilder b = new StringBuilder();
@@ -131,6 +146,9 @@ public String toString() {
131146
if (roles != null) {
132147
b.append(", roles=").append(roles);
133148
}
149+
if (attributes != null) {
150+
b.append(", attributes=").append(attributes);
151+
}
134152
return b.append(']').toString();
135153
}
136154

@@ -144,12 +162,13 @@ public boolean equals(Object obj) {
144162
&& Objects.equals(boundHosts, other.boundHosts)
145163
&& Objects.equals(name, other.name)
146164
&& Objects.equals(version, other.version)
147-
&& Objects.equals(roles, other.roles);
165+
&& Objects.equals(roles, other.roles)
166+
&& Objects.equals(attributes, other.attributes);
148167
}
149168

150169
@Override
151170
public int hashCode() {
152-
return Objects.hash(host, boundHosts, name, version, roles);
171+
return Objects.hash(host, boundHosts, name, version, roles, attributes);
153172
}
154173

155174
/**

‎client/rest/src/test/java/org/elasticsearch/client/HasAttributeNodeSelectorTests.java‎

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
/*
2+
* Licensed to Elasticsearch under one or more contributor
3+
* license agreements. See the NOTICE file distributed with
4+
* this work for additional information regarding copyright
5+
* ownership. Elasticsearch licenses this file to you under
6+
* the Apache License, Version 2.0 (the "License"); you may
7+
* not use this file except in compliance with the License.
8+
* You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
package org.elasticsearch.client;
21+
22+
import org.apache.http.HttpHost;
23+
import org.elasticsearch.client.Node.Roles;
24+
25+
import java.util.ArrayList;
26+
import java.util.Arrays;
27+
import java.util.Collections;
28+
import java.util.List;
29+
import java.util.Map;
30+
31+
import static java.util.Collections.singletonList;
32+
import static java.util.Collections.singletonMap;
33+
import static org.junit.Assert.assertEquals;
34+
35+
public class HasAttributeNodeSelectorTests extends RestClientTestCase {
36+
public void testHasAttribute() {
37+
Node hasAttributeValue = dummyNode(singletonMap("attr", singletonList("val")));
38+
Node hasAttributeButNotValue = dummyNode(singletonMap("attr", singletonList("notval")));
39+
Node hasAttributeValueInList = dummyNode(singletonMap("attr", Arrays.asList("val", "notval")));
40+
Node notHasAttribute = dummyNode(singletonMap("notattr", singletonList("val")));
41+
List<Node> nodes = new ArrayList<>();
42+
nodes.add(hasAttributeValue);
43+
nodes.add(hasAttributeButNotValue);
44+
nodes.add(hasAttributeValueInList);
45+
nodes.add(notHasAttribute);
46+
List<Node> expected = new ArrayList<>();
47+
expected.add(hasAttributeValue);
48+
expected.add(hasAttributeValueInList);
49+
new HasAttributeNodeSelector("attr", "val").select(nodes);
50+
assertEquals(expected, nodes);
51+
}
52+
53+
private static Node dummyNode(Map<String, List<String>> attributes) {
54+
return new Node(new HttpHost("dummy"), Collections.<HttpHost>emptySet(),
55+
randomAsciiAlphanumOfLength(5), randomAsciiAlphanumOfLength(5),
56+
new Roles(randomBoolean(), randomBoolean(), randomBoolean()),
57+
attributes);
58+
}
59+
}

‎client/rest/src/test/java/org/elasticsearch/client/NodeSelectorTests.java‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -63,9 +63,10 @@ public void testNotMasterOnly() {
6363
assertEquals(expected, nodes);
6464
}
6565

66-
private Node dummyNode(boolean master, boolean data, boolean ingest) {
66+
private static Node dummyNode(boolean master, boolean data, boolean ingest) {
6767
return new Node(new HttpHost("dummy"), Collections.<HttpHost>emptySet(),
6868
randomAsciiAlphanumOfLength(5), randomAsciiAlphanumOfLength(5),
69-
new Roles(master, data, ingest));
69+
new Roles(master, data, ingest),
70+
Collections.<String, List<String>>emptyMap());
7071
}
7172
}

‎client/rest/src/test/java/org/elasticsearch/client/NodeTests.java‎

Lines changed: 34 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -23,49 +23,67 @@
2323
import org.elasticsearch.client.Node.Roles;
2424

2525
import java.util.Arrays;
26+
import java.util.HashMap;
2627
import java.util.HashSet;
28+
import java.util.List;
29+
import java.util.Map;
2730

2831
import static java.util.Collections.singleton;
32+
import static java.util.Collections.singletonList;
33+
import static java.util.Collections.singletonMap;
2934
import static org.junit.Assert.assertEquals;
3035
import static org.junit.Assert.assertFalse;
3136
import static org.junit.Assert.assertTrue;
3237

3338
public class NodeTests extends RestClientTestCase {
3439
public void testToString() {
40+
Map<String, List<String>> attributes = new HashMap<>();
41+
attributes.put("foo", singletonList("bar"));
42+
attributes.put("baz", Arrays.asList("bort", "zoom"));
3543
assertEquals("[host=http://1]", new Node(new HttpHost("1")).toString());
44+
assertEquals("[host=http://1, attributes={foo=[bar], baz=[bort, zoom]}]",
45+
new Node(new HttpHost("1"), null, null, null, null, attributes).toString());
3646
assertEquals("[host=http://1, roles=mdi]", new Node(new HttpHost("1"),
37-
null, null, null, new Roles(true, true, true)).toString());
47+
null, null, null, new Roles(true, true, true), null).toString());
3848
assertEquals("[host=http://1, version=ver]", new Node(new HttpHost("1"),
39-
null, null, "ver", null).toString());
49+
null, null, "ver", null, null).toString());
4050
assertEquals("[host=http://1, name=nam]", new Node(new HttpHost("1"),
41-
null, "nam", null, null).toString());
51+
null, "nam", null, null, null).toString());
4252
assertEquals("[host=http://1, bound=[http://1, http://2]]", new Node(new HttpHost("1"),
43-
new HashSet<>(Arrays.asList(new HttpHost("1"), new HttpHost("2"))), null, null, null).toString());
44-
assertEquals("[host=http://1, bound=[http://1, http://2], name=nam, version=ver, roles=m]",
53+
new HashSet<>(Arrays.asList(new HttpHost("1"), new HttpHost("2"))), null, null, null, null).toString());
54+
assertEquals(
55+
"[host=http://1, bound=[http://1, http://2], name=nam, version=ver, roles=m, attributes={foo=[bar], baz=[bort, zoom]}]",
4556
new Node(new HttpHost("1"), new HashSet<>(Arrays.asList(new HttpHost("1"), new HttpHost("2"))),
46-
"nam", "ver", new Roles(true, false, false)).toString());
57+
"nam", "ver", new Roles(true, false, false), attributes).toString());
4758

4859
}
4960

5061
public void testEqualsAndHashCode() {
5162
HttpHost host = new HttpHost(randomAsciiAlphanumOfLength(5));
5263
Node node = new Node(host,
53-
randomBoolean() ? null : singleton(host),
54-
randomBoolean() ? null : randomAsciiAlphanumOfLength(5),
55-
randomBoolean() ? null : randomAsciiAlphanumOfLength(5),
56-
randomBoolean() ? null : new Roles(true, true, true));
64+
randomBoolean() ? null : singleton(host),
65+
randomBoolean() ? null : randomAsciiAlphanumOfLength(5),
66+
randomBoolean() ? null : randomAsciiAlphanumOfLength(5),
67+
randomBoolean() ? null : new Roles(true, true, true),
68+
randomBoolean() ? null : singletonMap("foo", singletonList("bar")));
5769
assertFalse(node.equals(null));
5870
assertTrue(node.equals(node));
5971
assertEquals(node.hashCode(), node.hashCode());
60-
Node copy = new Node(host, node.getBoundHosts(), node.getName(), node.getVersion(), node.getRoles());
72+
Node copy = new Node(host, node.getBoundHosts(), node.getName(), node.getVersion(),
73+
node.getRoles(), node.getAttributes());
6174
assertTrue(node.equals(copy));
6275
assertEquals(node.hashCode(), copy.hashCode());
6376
assertFalse(node.equals(new Node(new HttpHost(host.toHostString() + "changed"), node.getBoundHosts(),
64-
node.getName(), node.getVersion(), node.getRoles())));
77+
node.getName(), node.getVersion(), node.getRoles(), node.getAttributes())));
6578
assertFalse(node.equals(new Node(host, new HashSet<>(Arrays.asList(host, new HttpHost(host.toHostString() + "changed"))),
66-
node.getName(), node.getVersion(), node.getRoles())));
67-
assertFalse(node.equals(new Node(host, node.getBoundHosts(), node.getName() + "changed", node.getVersion(), node.getRoles())));
68-
assertFalse(node.equals(new Node(host, node.getBoundHosts(), node.getName(), node.getVersion() + "changed", node.getRoles())));
69-
assertFalse(node.equals(new Node(host, node.getBoundHosts(), node.getName(), node.getVersion(), new Roles(false, false, false))));
79+
node.getName(), node.getVersion(), node.getRoles(), node.getAttributes())));
80+
assertFalse(node.equals(new Node(host, node.getBoundHosts(), node.getName() + "changed",
81+
node.getVersion(), node.getRoles(), node.getAttributes())));
82+
assertFalse(node.equals(new Node(host, node.getBoundHosts(), node.getName(),
83+
node.getVersion() + "changed", node.getRoles(), node.getAttributes())));
84+
assertFalse(node.equals(new Node(host, node.getBoundHosts(), node.getName(),
85+
node.getVersion(), new Roles(false, false, false), node.getAttributes())));
86+
assertFalse(node.equals(new Node(host, node.getBoundHosts(), node.getName(),
87+
node.getVersion(), node.getRoles(), singletonMap("bort", singletonList("bing")))));
7088
}
7189
}

‎client/rest/src/test/java/org/elasticsearch/client/RestClientMultipleHostsTests.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -342,7 +342,7 @@ public void testSetNodes() throws IOException {
342342
List<Node> newNodes = new ArrayList<>(nodes.size());
343343
for (int i = 0; i < nodes.size(); i++) {
344344
Roles roles = i == 0 ? new Roles(false, true, true) : new Roles(true, false, false);
345-
newNodes.add(new Node(nodes.get(i).getHost(), null, null, null, roles));
345+
newNodes.add(new Node(nodes.get(i).getHost(), null, null, null, roles, null));
346346
}
347347
restClient.setNodes(newNodes);
348348
int rounds = between(1, 10);

‎client/rest/src/test/java/org/elasticsearch/client/RestClientTests.java‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -341,9 +341,9 @@ public void testNullPath() throws IOException {
341341
}
342342

343343
public void testSelectHosts() throws IOException {
344-
Node n1 = new Node(new HttpHost("1"), null, null, "1", null);
345-
Node n2 = new Node(new HttpHost("2"), null, null, "2", null);
346-
Node n3 = new Node(new HttpHost("3"), null, null, "3", null);
344+
Node n1 = new Node(new HttpHost("1"), null, null, "1", null, null);
345+
Node n2 = new Node(new HttpHost("2"), null, null, "2", null, null);
346+
Node n3 = new Node(new HttpHost("3"), null, null, "3", null, null);
347347

348348
NodeSelector not1 = new NodeSelector() {
349349
@Override

‎client/rest/src/test/java/org/elasticsearch/client/documentation/RestClientDocumentation.java‎

Lines changed: 15 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@
3636
import org.apache.http.ssl.SSLContextBuilder;
3737
import org.apache.http.ssl.SSLContexts;
3838
import org.apache.http.util.EntityUtils;
39+
import org.elasticsearch.client.HasAttributeNodeSelector;
3940
import org.elasticsearch.client.HttpAsyncResponseConsumerFactory.HeapBufferedResponseConsumerFactory;
4041
import org.elasticsearch.client.Node;
4142
import org.elasticsearch.client.NodeSelector;
@@ -190,11 +191,20 @@ public void onFailure(Exception exception) {
190191
//tag::rest-client-options-set-singleton
191192
request.setOptions(COMMON_OPTIONS);
192193
//end::rest-client-options-set-singleton
193-
//tag::rest-client-options-customize
194-
RequestOptions.Builder options = COMMON_OPTIONS.toBuilder();
195-
options.addHeader("cats", "knock things off of other things");
196-
request.setOptions(options);
197-
//end::rest-client-options-customize
194+
{
195+
//tag::rest-client-options-customize-header
196+
RequestOptions.Builder options = COMMON_OPTIONS.toBuilder();
197+
options.addHeader("cats", "knock things off of other things");
198+
request.setOptions(options);
199+
//end::rest-client-options-customize-header
200+
}
201+
{
202+
//tag::rest-client-options-customize-attribute
203+
RequestOptions.Builder options = COMMON_OPTIONS.toBuilder();
204+
options.setNodeSelector(new HasAttributeNodeSelector("rack", "c12")); // <1>
205+
request.setOptions(options);
206+
//end::rest-client-options-customize-attribute
207+
}
198208
}
199209
{
200210
HttpEntity[] documents = new HttpEntity[10];

0 commit comments

Comments
 (0)