Examples with KnownLength io.grpc.KnownLength used on opensource projects

Example 1 with KnownLength

use of io.grpc.KnownLength in project grpc-java by grpc.

the class ProtoLiteUtils method marshaller.

/** Create a {@code Marshaller} for protos of the same type as {@code defaultInstance}. */
public static <T extends MessageLite> Marshaller<T> marshaller(final T defaultInstance) {
    @SuppressWarnings("unchecked") final Parser<T> parser = (Parser<T>) defaultInstance.getParserForType();
    // TODO(ejona): consider changing return type to PrototypeMarshaller (assuming ABI safe)
    return new PrototypeMarshaller<T>() {

        @SuppressWarnings("unchecked")
        @Override
        public Class<T> getMessageClass() {
            // Precisely T since protobuf doesn't let messages extend other messages.
            return (Class<T>) defaultInstance.getClass();
        }

        @Override
        public T getMessagePrototype() {
            return defaultInstance;
        }

        @Override
        public InputStream stream(T value) {
            return new ProtoInputStream(value, parser);
        }

        @Override
        public T parse(InputStream stream) {
            if (stream instanceof ProtoInputStream) {
                ProtoInputStream protoStream = (ProtoInputStream) stream;
                // to enable this optimization.
                if (protoStream.parser() == parser) {
                    try {
                        @SuppressWarnings("unchecked") T message = (T) ((ProtoInputStream) stream).message();
                        return message;
                    } catch (IllegalStateException ex) {
                    // Stream must have been read from, which is a strange state. Since the point of this
                    // optimization is to be transparent, instead of throwing an error we'll continue,
                    // even though it seems likely there's a bug.
                    }
                }
            }
            CodedInputStream cis = null;
            try {
                if (stream instanceof KnownLength) {
                    int size = stream.available();
                    if (size > 0 && size <= GrpcUtil.DEFAULT_MAX_MESSAGE_SIZE) {
                        // buf should not be used after this method has returned.
                        byte[] buf = bufs.get().get();
                        if (buf == null || buf.length < size) {
                            buf = new byte[size];
                            bufs.set(new WeakReference<byte[]>(buf));
                        }
                        int chunkSize;
                        int position = 0;
                        while ((chunkSize = stream.read(buf, position, size - position)) != -1) {
                            position += chunkSize;
                        }
                        if (size != position) {
                            throw new RuntimeException("size inaccurate: " + size + " != " + position);
                        }
                        cis = CodedInputStream.newInstance(buf, 0, size);
                    } else if (size == 0) {
                        return defaultInstance;
                    }
                }
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
            if (cis == null) {
                cis = CodedInputStream.newInstance(stream);
            }
            // Pre-create the CodedInputStream so that we can remove the size limit restriction
            // when parsing.
            cis.setSizeLimit(Integer.MAX_VALUE);
            try {
                return parseFrom(cis);
            } catch (InvalidProtocolBufferException ipbe) {
                throw Status.INTERNAL.withDescription("Invalid protobuf byte sequence").withCause(ipbe).asRuntimeException();
            }
        }

        private T parseFrom(CodedInputStream stream) throws InvalidProtocolBufferException {
            T message = parser.parseFrom(stream, globalRegistry);
            try {
                stream.checkLastTagWas(0);
                return message;
            } catch (InvalidProtocolBufferException e) {
                e.setUnfinishedMessage(message);
                throw e;
            }
        }
    };
}