ejona | 7f23d5c | 2014-06-18 13:23:32 -0700 | [diff] [blame] | 1 | package com.google.net.stubby.newtransport; |
| 2 | |
| 3 | import com.google.common.io.ByteStreams; |
| 4 | import com.google.net.stubby.GrpcFramingUtil; |
| 5 | import com.google.net.stubby.Operation; |
| 6 | import com.google.net.stubby.Status; |
| 7 | import com.google.net.stubby.transport.Transport; |
| 8 | |
| 9 | import java.io.ByteArrayInputStream; |
| 10 | import java.io.DataInputStream; |
| 11 | import java.io.IOException; |
| 12 | import java.io.InputStream; |
| 13 | |
| 14 | /** |
| 15 | * Base implementation that joins a sequence of framed GRPC data produced by a {@link Framer}, |
| 16 | * reconstructs their messages and hands them off to a receiving {@link Operation} |
| 17 | */ |
| 18 | public abstract class Deframer<F> implements Framer.Sink<F> { |
| 19 | |
| 20 | /** |
| 21 | * Unset frame length |
| 22 | */ |
| 23 | private static final int LENGTH_NOT_SET = -1; |
| 24 | |
ejona | 9d50299 | 2014-09-22 12:23:19 -0700 | [diff] [blame^] | 25 | private final GrpcDeframer.Sink target; |
ejona | 7f23d5c | 2014-06-18 13:23:32 -0700 | [diff] [blame] | 26 | private boolean inFrame; |
| 27 | private byte currentFlags; |
| 28 | private int currentLength = LENGTH_NOT_SET; |
nathanmittler | cc7cdb1 | 2014-07-11 12:00:32 -0700 | [diff] [blame] | 29 | private boolean statusDelivered; |
ejona | 7f23d5c | 2014-06-18 13:23:32 -0700 | [diff] [blame] | 30 | |
ejona | 9d50299 | 2014-09-22 12:23:19 -0700 | [diff] [blame^] | 31 | public Deframer(GrpcDeframer.Sink target) { |
ejona | 7f23d5c | 2014-06-18 13:23:32 -0700 | [diff] [blame] | 32 | this.target = target; |
| 33 | } |
| 34 | |
| 35 | @Override |
| 36 | public void deliverFrame(F frame, boolean endOfStream) { |
| 37 | int remaining = internalDeliverFrame(frame); |
nathanmittler | cc7cdb1 | 2014-07-11 12:00:32 -0700 | [diff] [blame] | 38 | if (endOfStream) { |
| 39 | if (remaining > 0) { |
| 40 | writeStatus(new Status(Transport.Code.UNKNOWN, "EOF on incomplete frame")); |
| 41 | } else if (!statusDelivered) { |
| 42 | writeStatus(Status.OK); |
| 43 | } |
ejona | 7f23d5c | 2014-06-18 13:23:32 -0700 | [diff] [blame] | 44 | } |
| 45 | } |
| 46 | |
| 47 | /** |
| 48 | * Consume a frame of bytes provided by the transport. Note that transport framing is not |
| 49 | * aligned on GRPC frame boundaries so this code needs to do bounds checking and buffering |
| 50 | * across transport frame boundaries. |
| 51 | * |
| 52 | * @return the number of unconsumed bytes remaining in the buffer |
| 53 | */ |
| 54 | private int internalDeliverFrame(F frame) { |
| 55 | try { |
| 56 | frame = decompress(frame); |
| 57 | DataInputStream grpcStream = prefix(frame); |
| 58 | // Loop until no more GRPC frames can be fully decoded |
| 59 | while (true) { |
| 60 | if (!inFrame) { |
| 61 | // Not in frame so attempt to read flags |
| 62 | if (!ensure(grpcStream, GrpcFramingUtil.FRAME_TYPE_LENGTH)) { |
| 63 | return consolidate(); |
| 64 | } |
| 65 | currentFlags = grpcStream.readByte(); |
| 66 | inFrame = true; |
| 67 | } |
| 68 | if (currentLength == LENGTH_NOT_SET) { |
| 69 | // Read the frame length |
| 70 | if (!ensure(grpcStream, GrpcFramingUtil.FRAME_LENGTH)) { |
| 71 | return consolidate(); |
| 72 | } |
| 73 | currentLength = grpcStream.readInt(); |
| 74 | } |
| 75 | // Ensure that the entire frame length is available to read |
| 76 | InputStream framedChunk = ensureMessage(grpcStream, currentLength); |
| 77 | if (framedChunk == null) { |
| 78 | // Insufficient bytes available |
| 79 | return consolidate(); |
| 80 | } |
| 81 | if (GrpcFramingUtil.isPayloadFrame(currentFlags)) { |
| 82 | // Advance stream now, because target.addPayload() may not or may process the frame on |
| 83 | // another thread. |
| 84 | framedChunk = new ByteArrayInputStream(ByteStreams.toByteArray(framedChunk)); |
| 85 | try { |
| 86 | // Report payload to the receiving operation |
nathanmittler | 43d2fcc | 2014-08-29 21:58:54 -0700 | [diff] [blame] | 87 | target.messageRead(framedChunk, currentLength); |
ejona | 7f23d5c | 2014-06-18 13:23:32 -0700 | [diff] [blame] | 88 | } finally { |
| 89 | currentLength = LENGTH_NOT_SET; |
| 90 | inFrame = false; |
| 91 | } |
ejona | 7f23d5c | 2014-06-18 13:23:32 -0700 | [diff] [blame] | 92 | } else if (GrpcFramingUtil.isStatusFrame(currentFlags)) { |
| 93 | int status = framedChunk.read() << 8 | framedChunk.read(); |
| 94 | Transport.Code code = Transport.Code.valueOf(status); |
| 95 | // TODO(user): Resolve what to do with remainder of framedChunk |
| 96 | try { |
| 97 | if (code == null) { |
| 98 | // Log for unknown code |
nathanmittler | cc7cdb1 | 2014-07-11 12:00:32 -0700 | [diff] [blame] | 99 | writeStatus(new Status(Transport.Code.UNKNOWN, "Unknown status code " + status)); |
ejona | 7f23d5c | 2014-06-18 13:23:32 -0700 | [diff] [blame] | 100 | } else { |
nathanmittler | cc7cdb1 | 2014-07-11 12:00:32 -0700 | [diff] [blame] | 101 | writeStatus(new Status(code)); |
ejona | 7f23d5c | 2014-06-18 13:23:32 -0700 | [diff] [blame] | 102 | } |
| 103 | } finally { |
| 104 | currentLength = LENGTH_NOT_SET; |
| 105 | inFrame = false; |
| 106 | } |
| 107 | } |
| 108 | if (grpcStream.available() == 0) { |
| 109 | // We've processed all the data so consolidate the underlying buffers |
| 110 | return consolidate(); |
| 111 | } |
| 112 | } |
| 113 | } catch (IOException ioe) { |
| 114 | Status status = new Status(Transport.Code.UNKNOWN, ioe); |
nathanmittler | cc7cdb1 | 2014-07-11 12:00:32 -0700 | [diff] [blame] | 115 | writeStatus(status); |
ejona | 7f23d5c | 2014-06-18 13:23:32 -0700 | [diff] [blame] | 116 | throw status.asRuntimeException(); |
| 117 | } |
| 118 | } |
| 119 | |
| 120 | /** |
| 121 | * Return a stream view over the current buffer prefixed to the input frame |
| 122 | */ |
| 123 | protected abstract DataInputStream prefix(F frame) throws IOException; |
| 124 | |
| 125 | /** |
| 126 | * Consolidate the underlying buffers and return the number of buffered bytes remaining |
| 127 | */ |
| 128 | protected abstract int consolidate() throws IOException; |
| 129 | |
| 130 | /** |
| 131 | * Decompress the raw frame buffer prior to prefixing it. |
| 132 | */ |
| 133 | protected abstract F decompress(F frame) throws IOException; |
| 134 | |
| 135 | /** |
| 136 | * Ensure that {@code len} bytes are available in the buffer and frame |
| 137 | */ |
| 138 | private boolean ensure(InputStream input, int len) throws IOException { |
| 139 | return (input.available() >= len); |
| 140 | } |
| 141 | |
nathanmittler | cc7cdb1 | 2014-07-11 12:00:32 -0700 | [diff] [blame] | 142 | private void writeStatus(Status status) { |
ejona | 9d50299 | 2014-09-22 12:23:19 -0700 | [diff] [blame^] | 143 | target.statusRead(status); |
| 144 | target.endOfStream(); |
nathanmittler | cc7cdb1 | 2014-07-11 12:00:32 -0700 | [diff] [blame] | 145 | statusDelivered = true; |
| 146 | } |
| 147 | |
ejona | 7f23d5c | 2014-06-18 13:23:32 -0700 | [diff] [blame] | 148 | /** |
| 149 | * Return a message of {@code len} bytes than can be read from the buffer. If sufficient |
| 150 | * bytes are unavailable then buffer the available bytes and return null. |
| 151 | */ |
| 152 | private InputStream ensureMessage(InputStream input, int len) |
| 153 | throws IOException { |
| 154 | if (input.available() < len) { |
| 155 | return null; |
| 156 | } |
| 157 | return ByteStreams.limit(input, len); |
| 158 | } |
| 159 | } |