|
9 | 9 | package org.opensearch.transport; |
10 | 10 |
|
11 | 11 | import org.opensearch.common.bytes.ReleasableBytesReference; |
| 12 | +import org.opensearch.common.lease.Releasable; |
| 13 | +import org.opensearch.common.lease.Releasables; |
| 14 | +import org.opensearch.core.common.bytes.CompositeBytesReference; |
12 | 15 |
|
13 | | -import java.io.Closeable; |
14 | 16 | import java.io.IOException; |
| 17 | +import java.util.ArrayDeque; |
| 18 | +import java.util.ArrayList; |
15 | 19 | import java.util.function.BiConsumer; |
16 | 20 |
|
17 | 21 | /** |
18 | | - * Interface for handling inbound bytes. Can be implemented by different transport protocols. |
| 22 | + * Handler for inbound bytes, using {@link InboundDecoder} to decode headers |
| 23 | + * and {@link InboundAggregator} to assemble complete messages to forward to |
| 24 | + * the given message handler to parse the message payload. |
19 | 25 | */ |
20 | | -public interface InboundBytesHandler extends Closeable { |
| 26 | +class InboundBytesHandler { |
21 | 27 |
|
22 | | - public void doHandleBytes( |
23 | | - TcpChannel channel, |
24 | | - ReleasableBytesReference reference, |
25 | | - BiConsumer<TcpChannel, ProtocolInboundMessage> messageHandler |
26 | | - ) throws IOException; |
| 28 | + private static final ThreadLocal<ArrayList<Object>> fragmentList = ThreadLocal.withInitial(ArrayList::new); |
27 | 29 |
|
28 | | - public boolean canHandleBytes(ReleasableBytesReference reference); |
| 30 | + private final ArrayDeque<ReleasableBytesReference> pending; |
| 31 | + private final InboundDecoder decoder; |
| 32 | + private final InboundAggregator aggregator; |
| 33 | + private final StatsTracker statsTracker; |
| 34 | + private boolean isClosed = false; |
| 35 | + |
| 36 | + InboundBytesHandler( |
| 37 | + ArrayDeque<ReleasableBytesReference> pending, |
| 38 | + InboundDecoder decoder, |
| 39 | + InboundAggregator aggregator, |
| 40 | + StatsTracker statsTracker |
| 41 | + ) { |
| 42 | + this.pending = pending; |
| 43 | + this.decoder = decoder; |
| 44 | + this.aggregator = aggregator; |
| 45 | + this.statsTracker = statsTracker; |
| 46 | + } |
| 47 | + |
| 48 | + public void close() { |
| 49 | + isClosed = true; |
| 50 | + } |
| 51 | + |
| 52 | + public void doHandleBytes(TcpChannel channel, ReleasableBytesReference reference, BiConsumer<TcpChannel, InboundMessage> messageHandler) |
| 53 | + throws IOException { |
| 54 | + final ArrayList<Object> fragments = fragmentList.get(); |
| 55 | + boolean continueHandling = true; |
| 56 | + |
| 57 | + while (continueHandling && isClosed == false) { |
| 58 | + boolean continueDecoding = true; |
| 59 | + while (continueDecoding && pending.isEmpty() == false) { |
| 60 | + try (ReleasableBytesReference toDecode = getPendingBytes()) { |
| 61 | + final int bytesDecoded = decoder.decode(toDecode, fragments::add); |
| 62 | + if (bytesDecoded != 0) { |
| 63 | + releasePendingBytes(bytesDecoded); |
| 64 | + if (fragments.isEmpty() == false && endOfMessage(fragments.get(fragments.size() - 1))) { |
| 65 | + continueDecoding = false; |
| 66 | + } |
| 67 | + } else { |
| 68 | + continueDecoding = false; |
| 69 | + } |
| 70 | + } |
| 71 | + } |
| 72 | + |
| 73 | + if (fragments.isEmpty()) { |
| 74 | + continueHandling = false; |
| 75 | + } else { |
| 76 | + try { |
| 77 | + forwardFragments(channel, fragments, messageHandler); |
| 78 | + } finally { |
| 79 | + for (Object fragment : fragments) { |
| 80 | + if (fragment instanceof ReleasableBytesReference) { |
| 81 | + ((ReleasableBytesReference) fragment).close(); |
| 82 | + } |
| 83 | + } |
| 84 | + fragments.clear(); |
| 85 | + } |
| 86 | + } |
| 87 | + } |
| 88 | + } |
| 89 | + |
| 90 | + private ReleasableBytesReference getPendingBytes() { |
| 91 | + if (pending.size() == 1) { |
| 92 | + return pending.peekFirst().retain(); |
| 93 | + } else { |
| 94 | + final ReleasableBytesReference[] bytesReferences = new ReleasableBytesReference[pending.size()]; |
| 95 | + int index = 0; |
| 96 | + for (ReleasableBytesReference pendingReference : pending) { |
| 97 | + bytesReferences[index] = pendingReference.retain(); |
| 98 | + ++index; |
| 99 | + } |
| 100 | + final Releasable releasable = () -> Releasables.closeWhileHandlingException(bytesReferences); |
| 101 | + return new ReleasableBytesReference(CompositeBytesReference.of(bytesReferences), releasable); |
| 102 | + } |
| 103 | + } |
| 104 | + |
| 105 | + private void releasePendingBytes(int bytesConsumed) { |
| 106 | + int bytesToRelease = bytesConsumed; |
| 107 | + while (bytesToRelease != 0) { |
| 108 | + try (ReleasableBytesReference reference = pending.pollFirst()) { |
| 109 | + assert reference != null; |
| 110 | + if (bytesToRelease < reference.length()) { |
| 111 | + pending.addFirst(reference.retainedSlice(bytesToRelease, reference.length() - bytesToRelease)); |
| 112 | + bytesToRelease -= bytesToRelease; |
| 113 | + } else { |
| 114 | + bytesToRelease -= reference.length(); |
| 115 | + } |
| 116 | + } |
| 117 | + } |
| 118 | + } |
| 119 | + |
| 120 | + private boolean endOfMessage(Object fragment) { |
| 121 | + return fragment == InboundDecoder.PING || fragment == InboundDecoder.END_CONTENT || fragment instanceof Exception; |
| 122 | + } |
| 123 | + |
| 124 | + private void forwardFragments(TcpChannel channel, ArrayList<Object> fragments, BiConsumer<TcpChannel, InboundMessage> messageHandler) |
| 125 | + throws IOException { |
| 126 | + for (Object fragment : fragments) { |
| 127 | + if (fragment instanceof Header) { |
| 128 | + assert aggregator.isAggregating() == false; |
| 129 | + aggregator.headerReceived((Header) fragment); |
| 130 | + } else if (fragment == InboundDecoder.PING) { |
| 131 | + assert aggregator.isAggregating() == false; |
| 132 | + messageHandler.accept(channel, InboundMessage.PING); |
| 133 | + } else if (fragment == InboundDecoder.END_CONTENT) { |
| 134 | + assert aggregator.isAggregating(); |
| 135 | + try (InboundMessage aggregated = aggregator.finishAggregation()) { |
| 136 | + statsTracker.markMessageReceived(); |
| 137 | + messageHandler.accept(channel, aggregated); |
| 138 | + } |
| 139 | + } else { |
| 140 | + assert aggregator.isAggregating(); |
| 141 | + assert fragment instanceof ReleasableBytesReference; |
| 142 | + aggregator.aggregate((ReleasableBytesReference) fragment); |
| 143 | + } |
| 144 | + } |
| 145 | + } |
29 | 146 |
|
30 | | - @Override |
31 | | - void close(); |
32 | 147 | } |
0 commit comments