Skip to content

Commit ba25c9b

Browse files
niloc132niloc132
committed
Share code with the existing implementation
1 parent 1cc40d8 commit ba25c9b

7 files changed

Lines changed: 310 additions & 488 deletions

File tree

‎grpc-java/grpc-servlet-websocket-jakarta/src/main/java/io/grpc/servlet/web/websocket/AbstractWebSocketServerStream.java‎

Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,115 @@
1+
package io.grpc.servlet.web.websocket;
2+
3+
import io.grpc.Attributes;
4+
import io.grpc.InternalMetadata;
5+
import io.grpc.Metadata;
6+
import io.grpc.ServerStreamTracer;
7+
import io.grpc.internal.ServerTransportListener;
8+
import jakarta.websocket.Endpoint;
9+
import jakarta.websocket.EndpointConfig;
10+
import jakarta.websocket.Session;
11+
12+
import java.io.IOException;
13+
import java.io.UncheckedIOException;
14+
import java.nio.ByteBuffer;
15+
import java.nio.charset.StandardCharsets;
16+
import java.util.ArrayList;
17+
import java.util.Arrays;
18+
import java.util.List;
19+
20+
public abstract class AbstractWebSocketServerStream extends Endpoint {
21+
protected final ServerTransportListener transportListener;
22+
protected final List<? extends ServerStreamTracer.Factory> streamTracerFactories;
23+
protected final int maxInboundMessageSize;
24+
protected final Attributes attributes;
25+
26+
// assigned on open, always available
27+
protected Session websocketSession;
28+
29+
protected AbstractWebSocketServerStream(ServerTransportListener transportListener,
30+
List<? extends ServerStreamTracer.Factory> streamTracerFactories, int maxInboundMessageSize,
31+
Attributes attributes) {
32+
this.transportListener = transportListener;
33+
this.streamTracerFactories = streamTracerFactories;
34+
this.maxInboundMessageSize = maxInboundMessageSize;
35+
this.attributes = attributes;
36+
}
37+
38+
protected static Metadata readHeaders(ByteBuffer headerPayload) {
39+
// Headers are passed as ascii (browsers don't support binary), ":"-separated key/value pairs, separated on
40+
// "\r\n". The client implementation shows that values might be comma-separated, but we'll pass that through
41+
// directly as a plain string.
42+
//
43+
List<byte[]> byteArrays = new ArrayList<>();
44+
while (headerPayload.hasRemaining()) {
45+
int nameStart = headerPayload.position();
46+
while (headerPayload.hasRemaining() && headerPayload.get() != ':');
47+
int nameEnd = headerPayload.position() - 1;
48+
int valueStart = headerPayload.position() + 1;// assumes that the colon is followed by a space
49+
50+
while (headerPayload.hasRemaining() && headerPayload.get() != '\n');
51+
int valueEnd = headerPayload.position() - 2;// assumes that \n is preceded by a \r, this isnt generally
52+
// safe?
53+
if (valueEnd < valueStart) {
54+
valueEnd = valueStart;
55+
}
56+
int endOfLinePosition = headerPayload.position();
57+
58+
byte[] headerBytes = new byte[nameEnd - nameStart];
59+
headerPayload.position(nameStart);
60+
headerPayload.get(headerBytes);
61+
62+
byteArrays.add(headerBytes);
63+
if (Arrays.equals(headerBytes, "content-type".getBytes(StandardCharsets.US_ASCII))) {
64+
// rewrite grpc-web content type to matching grpc content type
65+
byteArrays.add("grpc+proto".getBytes(StandardCharsets.US_ASCII));
66+
// TODO support other formats like text, non-proto
67+
headerPayload.position(valueEnd);
68+
continue;
69+
}
70+
71+
// TODO check for binary header suffix
72+
// if (headerBytes.endsWith(Metadata.BINARY_HEADER_SUFFIX)) {
73+
//
74+
// } else {
75+
byte[] valueBytes = new byte[valueEnd - valueStart];
76+
headerPayload.position(valueStart);
77+
headerPayload.get(valueBytes);
78+
byteArrays.add(valueBytes);
79+
// }
80+
81+
headerPayload.position(endOfLinePosition);
82+
}
83+
84+
// add a te:trailers, as gRPC will expect it
85+
byteArrays.add("te".getBytes(StandardCharsets.US_ASCII));
86+
byteArrays.add("trailers".getBytes(StandardCharsets.US_ASCII));
87+
88+
// TODO to support text encoding
89+
90+
return InternalMetadata.newMetadata(byteArrays.toArray(new byte[][] {}));
91+
}
92+
93+
@Override
94+
public void onOpen(Session websocketSession, EndpointConfig config) {
95+
this.websocketSession = websocketSession;
96+
97+
websocketSession.addMessageHandler(String.class, this::onMessage);
98+
websocketSession.addMessageHandler(ByteBuffer.class, message -> {
99+
try {
100+
onMessage(message);
101+
} catch (IOException e) {
102+
throw new UncheckedIOException(e);
103+
}
104+
});
105+
106+
// Configure defaults present in some servlet containers to avoid some confusing limits. Subclasses
107+
// can override this method to control those defaults on their own.
108+
websocketSession.setMaxIdleTimeout(0);
109+
websocketSession.setMaxBinaryMessageBufferSize(Integer.MAX_VALUE);
110+
}
111+
112+
public abstract void onMessage(String message);
113+
114+
public abstract void onMessage(ByteBuffer message) throws IOException;
115+
}

‎grpc-java/grpc-servlet-websocket-jakarta/src/main/java/io/grpc/servlet/web/websocket/AbstractWebsocketStreamImpl.java‎

Lines changed: 133 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,133 @@
1+
package io.grpc.servlet.web.websocket;
2+
3+
import com.google.common.util.concurrent.MoreExecutors;
4+
import io.grpc.Attributes;
5+
import io.grpc.InternalLogId;
6+
import io.grpc.Metadata;
7+
import io.grpc.Status;
8+
import io.grpc.internal.AbstractServerStream;
9+
import io.grpc.internal.ReadableBuffer;
10+
import io.grpc.internal.SerializingExecutor;
11+
import io.grpc.internal.ServerTransportListener;
12+
import io.grpc.internal.StatsTraceContext;
13+
import io.grpc.internal.TransportTracer;
14+
import io.grpc.internal.WritableBufferAllocator;
15+
import jakarta.websocket.Session;
16+
17+
import java.io.IOException;
18+
import java.nio.ByteBuffer;
19+
import java.util.concurrent.CountDownLatch;
20+
import java.util.concurrent.TimeUnit;
21+
import java.util.logging.Level;
22+
import java.util.logging.Logger;
23+
24+
public abstract class AbstractWebsocketStreamImpl extends AbstractServerStream {
25+
public final class WebsocketTransportState extends TransportState {
26+
27+
private final SerializingExecutor transportThreadExecutor =
28+
new SerializingExecutor(MoreExecutors.directExecutor());
29+
private final Logger logger;
30+
31+
private WebsocketTransportState(int maxMessageSize, StatsTraceContext statsTraceCtx,
32+
TransportTracer transportTracer, Logger logger) {
33+
super(maxMessageSize, statsTraceCtx, transportTracer);
34+
this.logger = logger;
35+
}
36+
37+
@Override
38+
public void runOnTransportThread(Runnable r) {
39+
transportThreadExecutor.execute(r);
40+
}
41+
42+
@Override
43+
public void bytesRead(int numBytes) {
44+
// no-op, no flow-control yet
45+
}
46+
47+
@Override
48+
public void deframeFailed(Throwable cause) {
49+
if (logger.isLoggable(Level.FINE)) {
50+
logger.log(Level.FINE, String.format("[{%s}] Exception processing message", logId), cause);
51+
}
52+
cancel(Status.fromThrowable(cause));
53+
}
54+
}
55+
56+
protected final TransportState transportState;
57+
protected final Session websocketSession;
58+
protected final InternalLogId logId;
59+
protected final Attributes attributes;
60+
61+
public AbstractWebsocketStreamImpl(WritableBufferAllocator bufferAllocator, StatsTraceContext statsTraceCtx,
62+
int maxInboundMessageSize, Session websocketSession, InternalLogId logId, Attributes attributes,
63+
Logger logger) {
64+
super(bufferAllocator, statsTraceCtx);
65+
transportState =
66+
new WebsocketTransportState(maxInboundMessageSize, statsTraceCtx, new TransportTracer(), logger);
67+
this.websocketSession = websocketSession;
68+
this.logId = logId;
69+
this.attributes = attributes;
70+
}
71+
72+
protected static void writeAsciiHeadersToMessage(byte[][] serializedHeaders, ByteBuffer message) {
73+
for (int i = 0; i < serializedHeaders.length; i += 2) {
74+
message.put(serializedHeaders[i]);
75+
message.put((byte) ':');
76+
message.put((byte) ' ');
77+
message.put(serializedHeaders[i + 1]);
78+
message.put((byte) '\r');
79+
message.put((byte) '\n');
80+
}
81+
}
82+
83+
@Override
84+
public int streamId() {
85+
return -1;
86+
}
87+
88+
@Override
89+
public Attributes getAttributes() {
90+
return attributes;
91+
}
92+
93+
public void createStream(ServerTransportListener transportListener, String methodName, Metadata headers) {
94+
transportListener.streamCreated(this, methodName, headers);
95+
transportState().onStreamAllocated();
96+
}
97+
98+
public void inboundDataReceived(ReadableBuffer message, boolean endOfStream) {
99+
transportState().inboundDataReceived(message, endOfStream);
100+
}
101+
102+
public void transportReportStatus(Status status) {
103+
transportState().transportReportStatus(status);
104+
}
105+
106+
@Override
107+
public TransportState transportState() {
108+
return transportState;
109+
}
110+
111+
protected void cancelSink(Status status) {
112+
if (!websocketSession.isOpen() && Status.Code.DEADLINE_EXCEEDED == status.getCode()) {
113+
return;
114+
}
115+
transportState.runOnTransportThread(() -> transportState.transportReportStatus(status));
116+
// There is no way to RST_STREAM with CANCEL code, so write trailers instead
117+
close(Status.CANCELLED.withCause(status.asRuntimeException()), new Metadata());
118+
CountDownLatch countDownLatch = new CountDownLatch(1);
119+
transportState.runOnTransportThread(() -> {
120+
try {
121+
websocketSession.close();
122+
} catch (IOException ioException) {
123+
// already closing, ignore
124+
}
125+
countDownLatch.countDown();
126+
});
127+
try {
128+
countDownLatch.await(5, TimeUnit.SECONDS);
129+
} catch (InterruptedException e) {
130+
Thread.currentThread().interrupt();
131+
}
132+
}
133+
}

‎grpc-java/grpc-servlet-websocket-jakarta/src/main/java/io/grpc/servlet/web/websocket/ByteArrayWritableBuffer.java‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
package io.grpc.servlet.web.websocket;
2+
3+
import io.grpc.internal.WritableBuffer;
4+
5+
import static java.lang.Math.max;
6+
import static java.lang.Math.min;
7+
8+
final class ByteArrayWritableBuffer implements WritableBuffer {
9+
10+
private final int capacity;
11+
final byte[] bytes;
12+
private int index;
13+
14+
ByteArrayWritableBuffer(int capacityHint) {
15+
this.bytes = new byte[min(1024 * 1024, max(4096, capacityHint))];
16+
this.capacity = bytes.length;
17+
}
18+
19+
@Override
20+
public void write(byte[] src, int srcIndex, int length) {
21+
System.arraycopy(src, srcIndex, bytes, index, length);
22+
index += length;
23+
}
24+
25+
@Override
26+
public void write(byte b) {
27+
bytes[index++] = b;
28+
}
29+
30+
@Override
31+
public int writableBytes() {
32+
return capacity - index;
33+
}
34+
35+
@Override
36+
public int readableBytes() {
37+
return index;
38+
}
39+
40+
@Override
41+
public void release() {}
42+
}

0 commit comments

Comments
 (0)