| package com.google.net.stubby.transport; |
| |
| import com.google.net.stubby.GrpcFramingUtil; |
| import com.google.net.stubby.Status; |
| import com.google.protobuf.CodedOutputStream; |
| import com.google.protobuf.WireFormat; |
| |
| import java.io.IOException; |
| import java.io.InputStream; |
| import java.nio.ByteBuffer; |
| import java.nio.charset.Charset; |
| import java.util.Arrays; |
| |
| /** |
| * Default {@link Framer} implementation. |
| */ |
| public class MessageFramer implements Framer { |
| |
| /** |
| * Size of the GRPC message frame header which consists of |
| * 1 byte for the type (payload, context, status) |
| * 4 bytes for the length of the message |
| */ |
| private static final int MESSAGE_HEADER_SIZE = 5; |
| |
| /** |
| * UTF-8 charset which is used for key name encoding in context values |
| */ |
| private static final Charset UTF_8 = Charset.forName("UTF-8"); |
| |
| /** |
| * Precomputed protobuf tags for ContextValue |
| */ |
| private static final byte[] VALUE_TAG; |
| private static final byte[] KEY_TAG; |
| |
| |
| static { |
| // Initialize constants for serializing context-value in a protobuf compatible manner |
| try { |
| byte[] buf = new byte[8]; |
| CodedOutputStream coded = CodedOutputStream.newInstance(buf); |
| coded.writeTag(Transport.ContextValue.KEY_FIELD_NUMBER, WireFormat.WIRETYPE_LENGTH_DELIMITED); |
| coded.flush(); |
| KEY_TAG = Arrays.copyOf(buf, coded.getTotalBytesWritten()); |
| coded = CodedOutputStream.newInstance(buf); |
| coded.writeTag(Transport.ContextValue.VALUE_FIELD_NUMBER, |
| WireFormat.WIRETYPE_LENGTH_DELIMITED); |
| coded.flush(); |
| VALUE_TAG = Arrays.copyOf(buf, coded.getTotalBytesWritten()); |
| } catch (IOException ioe) { |
| // Unrecoverable |
| throw new RuntimeException(ioe); |
| } |
| } |
| |
| private CompressionFramer framer; |
| private final ByteBuffer scratch = ByteBuffer.allocate(16); |
| |
| public MessageFramer(int maxFrameSize) { |
| // TODO(user): maxFrameSize should probably come from a 'Platform' class |
| framer = new CompressionFramer(maxFrameSize, false, maxFrameSize / 16); |
| } |
| |
| /** |
| * Sets whether compression is encouraged. |
| */ |
| public void setAllowCompression(boolean enable) { |
| framer.setAllowCompression(enable); |
| } |
| |
| /** |
| * Set the preferred compression level for when compression is enabled. |
| * @param level the preferred compression level, or {@code -1} to use the framing default |
| * @see java.util.zip.Deflater#setLevel |
| */ |
| public void setCompressionLevel(int level) { |
| framer.setCompressionLevel(level); |
| } |
| |
| @Override |
| public void writePayload(InputStream message, boolean flush, Sink sink) { |
| try { |
| scratch.clear(); |
| scratch.put(GrpcFramingUtil.PAYLOAD_FRAME); |
| int messageLength = message.available(); |
| scratch.putInt(messageLength); |
| framer.write(scratch.array(), 0, scratch.position(), sink); |
| if (messageLength != framer.write(message, sink)) { |
| throw new RuntimeException("InputStream's available() was inaccurate"); |
| } |
| framer.endOfMessage(sink); |
| if (flush && framer != null) { |
| framer.flush(sink); |
| } |
| } catch (IOException ioe) { |
| throw new RuntimeException(ioe); |
| } |
| } |
| |
| |
| @Override |
| public void writeContext(String key, InputStream message, boolean flush, Sink sink) { |
| try { |
| scratch.clear(); |
| scratch.put(GrpcFramingUtil.CONTEXT_VALUE_FRAME); |
| byte[] keyBytes = key.getBytes(UTF_8); |
| int lenKeyPrefix = KEY_TAG.length + |
| CodedOutputStream.computeRawVarint32Size(keyBytes.length); |
| int messageLen = message.available(); |
| int lenValPrefix = VALUE_TAG.length + CodedOutputStream.computeRawVarint32Size(messageLen); |
| int totalLen = lenKeyPrefix + keyBytes.length + lenValPrefix + messageLen; |
| scratch.putInt(totalLen); |
| framer.write(scratch.array(), 0, scratch.position(), sink); |
| |
| // Write key |
| scratch.clear(); |
| scratch.put(KEY_TAG); |
| writeRawVarInt32(keyBytes.length, scratch); |
| framer.write(scratch.array(), 0, scratch.position(), sink); |
| framer.write(keyBytes, 0, keyBytes.length, sink); |
| |
| // Write value |
| scratch.clear(); |
| scratch.put(VALUE_TAG); |
| writeRawVarInt32(messageLen, scratch); |
| framer.write(scratch.array(), 0, scratch.position(), sink); |
| if (messageLen != framer.write(message, sink)) { |
| throw new RuntimeException("InputStream's available() was inaccurate"); |
| } |
| framer.endOfMessage(sink); |
| if (flush && framer != null) { |
| framer.flush(sink); |
| } |
| } catch (IOException ioe) { |
| throw new RuntimeException(ioe); |
| } |
| } |
| |
| @Override |
| public void writeStatus(Status status, boolean flush, Sink sink) { |
| short code = (short) status.getCode().getNumber(); |
| scratch.clear(); |
| scratch.put(GrpcFramingUtil.STATUS_FRAME); |
| int length = 2; |
| scratch.putInt(length); |
| scratch.putShort(code); |
| framer.write(scratch.array(), 0, scratch.position(), sink); |
| framer.endOfMessage(sink); |
| if (flush && framer != null) { |
| framer.flush(sink); |
| } |
| } |
| |
| @Override |
| public void flush(Sink sink) { |
| framer.flush(sink); |
| } |
| |
| @Override |
| public void close() { |
| // TODO(user): Returning buffer to a pool would go here |
| framer = null; |
| } |
| |
| /** |
| * Write a raw VarInt32 to the buffer |
| */ |
| private static void writeRawVarInt32(int value, ByteBuffer bytebuf) { |
| while (true) { |
| if ((value & ~0x7F) == 0) { |
| bytebuf.put((byte) value); |
| return; |
| } else { |
| bytebuf.put((byte) ((value & 0x7F) | 0x80)); |
| value >>>= 7; |
| } |
| } |
| } |
| } |