Skip to content

Commit 733931b

Browse files
retareta
authored
Fix X-Opaque-Id header propagation (along with other response headers) for streaming Reactor Netty 4 transport (opensearch-project#20371)
* Fix X-Opaque-Id header propagation (along with other response headers) for streaming Reactor Netty 4 transport Signed-off-by: Andriy Redko <drreta@gmail.com> * Address code review comments Signed-off-by: Andriy Redko <drreta@gmail.com> --------- Signed-off-by: Andriy Redko <drreta@gmail.com>
1 parent 62eed1e commit 733931b

7 files changed

Lines changed: 231 additions & 28 deletions

File tree

‎CHANGELOG.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
2929
- Fix snapshot restore when an index sort is present ([#20284](https://github.com/opensearch-project/OpenSearch/pull/20284))
3030
- Fix SearchPhaseExecutionException to properly initCause ([#20320](https://github.com/opensearch-project/OpenSearch/pull/20320))
3131
- Fix `cluster.remote.<cluster_alias>.server_name` setting no populating SNI ([#20321](https://github.com/opensearch-project/OpenSearch/pull/20321))
32+
- Fix X-Opaque-Id header propagation (along with other response headers) for streaming Reactor Netty 4 transport ([#20371](https://github.com/opensearch-project/OpenSearch/pull/20371))
3233

3334
### Dependencies
3435
- Bump `com.google.auth:google-auth-library-oauth2-http` from 1.38.0 to 1.41.0 ([#20183](https://github.com/opensearch-project/OpenSearch/pull/20183))

‎client/rest/src/main/java/org/opensearch/client/StreamingRequest.java‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,30 @@ public Map<String, String> getParameters() {
7777
}
7878
}
7979

80+
/**
81+
* Set the portion of an HTTP request to OpenSearch that can be
82+
* manipulated without changing OpenSearch's behavior.
83+
*
84+
* @param options the options to be set.
85+
* @throws NullPointerException if {@code options} is null.
86+
*/
87+
public void setOptions(RequestOptions options) {
88+
Objects.requireNonNull(options, "options cannot be null");
89+
this.options = options;
90+
}
91+
92+
/**
93+
* Set the portion of an HTTP request to OpenSearch that can be
94+
* manipulated without changing OpenSearch's behavior.
95+
*
96+
* @param options the options to be set.
97+
* @throws NullPointerException if {@code options} is null.
98+
*/
99+
public void setOptions(RequestOptions.Builder options) {
100+
Objects.requireNonNull(options, "options cannot be null");
101+
this.options = options.build();
102+
}
103+
80104
/**
81105
* Add a query string parameter.
82106
* @param name the name of the url parameter. Must not be null.

‎client/rest/src/main/java/org/opensearch/client/StreamingResponse.java‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88

99
package org.opensearch.client;
1010

11+
import org.apache.hc.core5.http.Header;
1112
import org.apache.hc.core5.http.HttpHost;
1213
import org.apache.hc.core5.http.HttpResponse;
1314
import org.apache.hc.core5.http.Message;
@@ -93,4 +94,29 @@ public List<String> getWarnings() {
9394
.block()
9495
);
9596
}
97+
98+
/**
99+
* Returns a list of all headers returned in the response.
100+
*/
101+
public Header[] getHeaders() {
102+
return publisher.map(Message::getHead)
103+
.onErrorResume(ResponseException.class, e -> Mono.just(e.getResponse().getHttpResponse()))
104+
.map(HttpResponse::getHeaders)
105+
.block();
106+
}
107+
108+
/**
109+
* Returns the value of the first header with a specified name of this message.
110+
* If there is more than one matching header in the message the first element is returned.
111+
* If there is no matching header in the message <code>null</code> is returned.
112+
*
113+
* @param name header name
114+
*/
115+
public String getHeader(String name) {
116+
return publisher.map(Message::getHead)
117+
.onErrorResume(ResponseException.class, e -> Mono.just(e.getResponse().getHttpResponse()))
118+
.mapNotNull(response -> response.getFirstHeader(name))
119+
.mapNotNull(header -> header.getValue())
120+
.block();
121+
}
96122
}

‎plugins/transport-reactor-netty4/src/javaRestTest/java/org/opensearch/rest/ReactorNetty4StreamingIT.java‎

Lines changed: 115 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
package org.opensearch.rest;
1010

1111
import org.opensearch.client.Request;
12+
import org.opensearch.client.RequestOptions;
1213
import org.opensearch.client.Response;
1314
import org.opensearch.client.ResponseException;
1415
import org.opensearch.client.StreamingRequest;
@@ -48,8 +49,10 @@ public void tearDown() throws Exception {
4849
public void testStreamingRequestNoBatching() throws IOException {
4950
final VirtualTimeScheduler scheduler = VirtualTimeScheduler.create(true);
5051

51-
final Stream<String> stream = IntStream.range(1, 6)
52-
.mapToObj(id -> "{ \"index\": { \"_index\": \"test-streaming\", \"_id\": \"" + id + "\" } }\n" + "{ \"name\": \"josh\" }\n");
52+
final Stream<String> stream = IntStream.range(1, 6).mapToObj(id -> String.format(Locale.ENGLISH, """
53+
{ "index": { "_index": "test-streaming", "_id": "%d" } }
54+
{ "name": "josh" }
55+
""", id));
5356

5457
final Duration delay = Duration.ofMillis(1);
5558
final StreamingRequest<ByteBuffer> streamingRequest = new StreamingRequest<>(
@@ -86,9 +89,87 @@ public void testStreamingRequestNoBatching() throws IOException {
8689
assertThat(count, equalTo(5));
8790
}
8891

92+
public void testStreamingRequestOpaqueId() throws IOException {
93+
final VirtualTimeScheduler scheduler = VirtualTimeScheduler.create(true);
94+
95+
final Stream<String> stream = IntStream.range(1, 6).mapToObj(id -> String.format(Locale.ENGLISH, """
96+
{ "index": { "_index": "test-streaming", "_id": "%d" } }
97+
{ "name": "josh" }
98+
""", id));
99+
100+
final RequestOptions options = RequestOptions.DEFAULT.toBuilder().addHeader("X-Opaque-Id", "1").build();
101+
102+
final Duration delay = Duration.ofMillis(1);
103+
final StreamingRequest<ByteBuffer> streamingRequest = new StreamingRequest<>(
104+
"POST",
105+
"/_bulk/stream",
106+
Flux.fromStream(stream).delayElements(delay, scheduler).map(s -> ByteBuffer.wrap(s.getBytes(StandardCharsets.UTF_8)))
107+
);
108+
streamingRequest.addParameter("refresh", "true");
109+
streamingRequest.setOptions(options);
110+
111+
final StreamingResponse<ByteBuffer> streamingResponse = client().streamRequest(streamingRequest);
112+
scheduler.advanceTimeBy(delay); /* emit first element */
113+
114+
StepVerifier.create(Flux.from(streamingResponse.getBody()).map(b -> new String(b.array(), StandardCharsets.UTF_8)))
115+
.expectNextMatches(s -> s.contains("\"result\":\"created\"") && s.contains("\"_id\":\"1\""))
116+
.then(() -> scheduler.advanceTimeBy(delay))
117+
.expectNextMatches(s -> s.contains("\"result\":\"created\"") && s.contains("\"_id\":\"2\""))
118+
.then(() -> scheduler.advanceTimeBy(delay))
119+
.expectNextMatches(s -> s.contains("\"result\":\"created\"") && s.contains("\"_id\":\"3\""))
120+
.then(() -> scheduler.advanceTimeBy(delay))
121+
.expectNextMatches(s -> s.contains("\"result\":\"created\"") && s.contains("\"_id\":\"4\""))
122+
.then(() -> scheduler.advanceTimeBy(delay))
123+
.expectNextMatches(s -> s.contains("\"result\":\"created\"") && s.contains("\"_id\":\"5\""))
124+
.then(() -> scheduler.advanceTimeBy(delay))
125+
.expectComplete()
126+
.verify();
127+
128+
assertThat(streamingResponse.getHeader("X-Opaque-Id"), equalTo("1"));
129+
assertThat(streamingResponse.getStatusLine().getStatusCode(), equalTo(200));
130+
assertThat(streamingResponse.getWarnings(), empty());
131+
132+
final Request request = new Request("GET", "/test-streaming/_count");
133+
final Response response = client().performRequest(request);
134+
final ObjectPath objectPath = ObjectPath.createFromResponse(response);
135+
final Integer count = objectPath.evaluate("count");
136+
assertThat(count, equalTo(5));
137+
}
138+
139+
public void testStreamingRequestOpaqueIdTwice() throws IOException {
140+
final Stream<String> stream = IntStream.range(1, 6).mapToObj(id -> String.format(Locale.ENGLISH, """
141+
{ "index": { "_index": "test-streaming", "_id": "%d" } }
142+
{ "name": "josh" }
143+
""", id));
144+
145+
final RequestOptions options = RequestOptions.DEFAULT.toBuilder()
146+
.addHeader("X-Opaque-Id", "1")
147+
.addHeader("X-Opaque-Id", "2")
148+
.build();
149+
150+
final StreamingRequest<ByteBuffer> streamingRequest = new StreamingRequest<>(
151+
"POST",
152+
"/_bulk/stream",
153+
Flux.fromStream(stream).map(s -> ByteBuffer.wrap(s.getBytes(StandardCharsets.UTF_8)))
154+
);
155+
streamingRequest.addParameter("refresh", "true");
156+
streamingRequest.setOptions(options);
157+
158+
final StreamingResponse<ByteBuffer> streamingResponse = client().streamRequest(streamingRequest);
159+
160+
StepVerifier.create(Flux.from(streamingResponse.getBody()).map(b -> new String(b.array(), StandardCharsets.UTF_8)))
161+
.expectError(ResponseException.class)
162+
.verify();
163+
164+
assertThat(streamingResponse.getStatusLine().getStatusCode(), equalTo(400));
165+
assertThat(streamingResponse.getWarnings(), empty());
166+
}
167+
89168
public void testStreamingRequestOneBatchBySize() throws IOException, InterruptedException {
90-
final Stream<String> stream = IntStream.range(1, 6)
91-
.mapToObj(id -> "{ \"index\": { \"_index\": \"test-streaming\", \"_id\": \"" + id + "\" } }\n" + "{ \"name\": \"josh\" }\n");
169+
final Stream<String> stream = IntStream.range(1, 6).mapToObj(id -> String.format(Locale.ENGLISH, """
170+
{ "index": { "_index": "test-streaming", "_id": "%d" } }
171+
{ "name": "josh" }
172+
""", id));
92173

93174
final Duration delay = Duration.ofMillis(1);
94175
final StreamingRequest<ByteBuffer> streamingRequest = new StreamingRequest<>(
@@ -128,8 +209,10 @@ public void testStreamingRequestOneBatchBySize() throws IOException, Interrupted
128209
}
129210

130211
public void testStreamingRequestManyBatchesBySize() throws IOException {
131-
final Stream<String> stream = IntStream.range(1, 6)
132-
.mapToObj(id -> "{ \"index\": { \"_index\": \"test-streaming\", \"_id\": \"" + id + "\" } }\n" + "{ \"name\": \"josh\" }\n");
212+
final Stream<String> stream = IntStream.range(1, 6).mapToObj(id -> String.format(Locale.ENGLISH, """
213+
{ "index": { "_index": "test-streaming", "_id": "%d" } }
214+
{ "name": "josh" }
215+
""", id));
133216

134217
final Duration delay = Duration.ofMillis(1);
135218
final StreamingRequest<ByteBuffer> streamingRequest = new StreamingRequest<>(
@@ -171,8 +254,10 @@ public void testStreamingRequestManyBatchesBySize() throws IOException {
171254
}
172255

173256
public void testStreamingRequestManyBatchesByInterval() throws IOException {
174-
final Stream<String> stream = IntStream.range(1, 6)
175-
.mapToObj(id -> "{ \"index\": { \"_index\": \"test-streaming\", \"_id\": \"" + id + "\" } }\n" + "{ \"name\": \"josh\" }\n");
257+
final Stream<String> stream = IntStream.range(1, 6).mapToObj(id -> String.format(Locale.ENGLISH, """
258+
{ "index": { "_index": "test-streaming", "_id": "%d" } }
259+
{ "name": "josh" }
260+
""", id));
176261

177262
final Duration delay = Duration.ofMillis(500);
178263
final StreamingRequest<ByteBuffer> streamingRequest = new StreamingRequest<>(
@@ -214,8 +299,10 @@ public void testStreamingRequestManyBatchesByInterval() throws IOException {
214299
}
215300

216301
public void testStreamingRequestManyBatchesByIntervalAndSize() throws IOException {
217-
final Stream<String> stream = IntStream.range(1, 6)
218-
.mapToObj(id -> "{ \"index\": { \"_index\": \"test-streaming\", \"_id\": \"" + id + "\" } }\n" + "{ \"name\": \"josh\" }\n");
302+
final Stream<String> stream = IntStream.range(1, 6).mapToObj(id -> String.format(Locale.ENGLISH, """
303+
{ "index": { "_index": "test-streaming", "_id": "%d" } }
304+
{ "name": "josh" }
305+
""", id));
219306

220307
final Duration delay = Duration.ofSeconds(1);
221308
final StreamingRequest<ByteBuffer> streamingRequest = new StreamingRequest<>(
@@ -248,9 +335,10 @@ public void testStreamingRequestManyBatchesByIntervalAndSize() throws IOExceptio
248335
}
249336

250337
public void testStreamingBadRequest() throws IOException {
251-
final Stream<String> stream = Stream.of(
252-
"{ \"index\": { \"_index\": \"test-streaming\", \"_id\": \"1\" } }\n" + "{ \"name\": \"josh\" }\n"
253-
);
338+
final Stream<String> stream = Stream.of("""
339+
{ "index": { "_index": "test-streaming", "_id": "1" } }
340+
{ "name": "josh" }
341+
""");
254342

255343
final StreamingRequest<ByteBuffer> streamingRequest = new StreamingRequest<>(
256344
"POST",
@@ -272,10 +360,12 @@ public void testStreamingBadRequest() throws IOException {
272360
public void testStreamingBadStream() throws IOException {
273361
final VirtualTimeScheduler scheduler = VirtualTimeScheduler.create(true);
274362

275-
final Stream<String> stream = Stream.of(
276-
"{ \"index\": { \"_index\": \"test-streaming\", \"_id\": \"1\" } }\n" + "{ \"name\": \"josh\" }\n",
277-
"{ \"name\": \"josh\" }\n"
278-
);
363+
final Stream<String> stream = Stream.of("""
364+
{ "index": { "_index": "test-streaming", "_id": "1" } }
365+
{ "name": "josh" }
366+
""", """
367+
{ "name": "josh" }
368+
""");
279369

280370
final Duration delay = Duration.ofMillis(1);
281371
final StreamingRequest<ByteBuffer> streamingRequest = new StreamingRequest<>(
@@ -300,13 +390,10 @@ public void testStreamingBadStream() throws IOException {
300390
}
301391

302392
public void testStreamingLargeDocument() throws IOException {
303-
final Stream<String> stream = Stream.of(
304-
String.format(
305-
Locale.getDefault(),
306-
"{ \"index\": { \"_index\": \"test-streaming\", \"_id\": \"1\" } }\n{ \"name\": \"%s\" }\n",
307-
randomAlphaOfLength(7000)
308-
)
309-
);
393+
final Stream<String> stream = Stream.of(String.format(Locale.getDefault(), """
394+
{ "index": { "_index": "test-streaming", "_id": "1" } }
395+
{ "name": "%s" }
396+
""", randomAlphaOfLength(7000)));
310397

311398
final StreamingRequest<ByteBuffer> streamingRequest = new StreamingRequest<>(
312399
"POST",
@@ -329,7 +416,10 @@ public void testStreamingLargeDocumentThatExceedsChunkSize() throws IOException
329416
final Stream<String> stream = Stream.of(
330417
String.format(
331418
Locale.getDefault(),
332-
"{ \"index\": { \"_index\": \"test-streaming\", \"_id\": \"1\" } }\n{ \"name\": \"%s\" }\n",
419+
"""
420+
{ "index": { "_index": "test-streaming", "_id": "1" } }
421+
{ "name": "%s" }
422+
""",
333423
randomAlphaOfLength(9000) /* the default chunk size limit is set 8k */
334424
)
335425
);

‎plugins/transport-reactor-netty4/src/main/java/org/opensearch/http/reactor/netty4/ReactorNetty4StreamingHttpChannel.java‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -82,8 +82,10 @@ public void sendResponse(HttpResponse response, ActionListener<Void> listener) {
8282

8383
@Override
8484
public void prepareResponse(int status, Map<String, List<String>> headers) {
85-
this.response.status(status);
86-
headers.forEach((k, vs) -> vs.forEach(v -> this.response.addHeader(k, v)));
85+
if (this.response.hasSentHeaders() == false) {
86+
this.response.status(status);
87+
headers.forEach((k, vs) -> vs.forEach(v -> this.response.addHeader(k, v)));
88+
}
8789
}
8890

8991
@Override

‎plugins/transport-reactor-netty4/src/main/java/org/opensearch/http/reactor/netty4/ReactorNetty4StreamingResponseProducer.java‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,15 +10,29 @@
1010

1111
import org.opensearch.core.action.ActionListener;
1212

13+
import java.util.Queue;
14+
import java.util.concurrent.ConcurrentLinkedQueue;
15+
1316
import io.netty.handler.codec.http.HttpContent;
1417
import org.reactivestreams.Publisher;
1518
import org.reactivestreams.Subscriber;
1619
import reactor.core.publisher.Flux;
1720
import reactor.core.publisher.FluxSink;
1821

1922
class ReactorNetty4StreamingResponseProducer implements StreamingHttpContentSender, Publisher<HttpContent> {
23+
// Buffer up to 64 messages, otherwise fail with buffer overflow (IllegalStateException)
24+
private static final int BUFFERED_QUEUE_SIZE = 64;
25+
2026
private final Publisher<HttpContent> sender;
2127
private volatile FluxSink<HttpContent> emitter;
28+
// We choose unbounded queue as the safe buffer here (since it does not support capacity bounds),
29+
// realistically we should only see a single final (last content) response being deferred. The BUFFERED_QUEUE_SIZE
30+
// check takes care of potential overflows.
31+
private final Queue<DelayedHttpContent> queue = new ConcurrentLinkedQueue<>();
32+
33+
// Holds the {@code HttpContent} for deferred delivery
34+
private record DelayedHttpContent(HttpContent content, ActionListener<Void> listener, boolean isLast) {
35+
};
2236

2337
ReactorNetty4StreamingResponseProducer() {
2438
this.sender = Flux.create(emitter -> register(emitter));
@@ -30,6 +44,20 @@ private void register(FluxSink<HttpContent> emitter) {
3044

3145
@Override
3246
public void send(HttpContent content, ActionListener<Void> listener, boolean isLast) {
47+
// In case when the exception triggers the response **before** the emitter is being
48+
// created, we defer the send till the subscribe call happens.
49+
if (isReady() == false) {
50+
queue.offer(new DelayedHttpContent(content, listener, isLast));
51+
if (queue.size() > BUFFERED_QUEUE_SIZE) {
52+
final IllegalStateException ex = new IllegalStateException(
53+
"The buffered queue size limit is exceeded: " + BUFFERED_QUEUE_SIZE
54+
);
55+
listener.onFailure(ex);
56+
throw ex;
57+
}
58+
return;
59+
}
60+
3361
try {
3462
emitter.next(content);
3563
listener.onResponse(null);
@@ -45,6 +73,11 @@ public void send(HttpContent content, ActionListener<Void> listener, boolean isL
4573
@Override
4674
public void subscribe(Subscriber<? super HttpContent> s) {
4775
sender.subscribe(s);
76+
77+
DelayedHttpContent content = null;
78+
while ((content = queue.poll()) != null) {
79+
send(content.content(), content.listener(), content.isLast());
80+
}
4881
}
4982

5083
@Override

0 commit comments

Comments
 (0)