org.apache.thrift.protocol.TProtocol Java Examples

@SuppressWarnings("rawtypes")
private void writeResult(final TProtocol out, final TMessage msg, final WriterHandler onComplete, TBase args,
		final TBase result) {
	try {
		onComplete.beforeWrite(msg, args, result);
		// if (!isOneway()) {
		out.writeMessageBegin(new TMessage(msg.name, TMessageType.REPLY, msg.seqid));
		if (result != null) {
			result.write(out);
		} else {
			out.writeStructBegin(null);
			out.writeFieldStop();
			out.writeStructEnd();
		}
		out.writeMessageEnd();
		out.getTransport().flush();
		// }
		onComplete.afterWrite(msg, null, TMessageType.REPLY, args, result);
	} catch (Throwable e) {
		onComplete.afterWrite(msg, e, TMessageType.EXCEPTION, args, result);
	}
}
 

@Test
public void testFixedResponse() throws Exception {
  ThriftMockServer server = new ThriftMockServer(9999);
  Thread t = new Thread(server::start);
  t.start();

  server.setExpectReturn("sayHello", expectHelloResponse);

  Request request = new Request();
  request.setMsg("hello, i'm guest");
  TProtocol protocol = TProtocolUtil.initTProtocol("127.0.0.1", 9999);
  HelloService.Iface helloService = new HelloService.Client(protocol);
  Response  actualResponse = helloService.sayHello(request);
  LOG.info("actualResponse: " + actualResponse);
  Assert.assertTrue(actualResponse.getCode() == 200);
  server.stop();
}
 

@Override
public Pair<THBaseService.Client, TTransport> getClient() throws IOException {
  Preconditions.checkArgument(connection.getHost().startsWith("http"),
      "http client host must start with http or https");
  String url = connection.getHost() + ":" + connection.getPort();
  try {
    THttpClient httpClient = new THttpClient(url, connection.getHttpClient());
    for (Map.Entry<String, String> header : customHeader.entrySet()) {
      httpClient.setCustomHeader(header.getKey(), header.getValue());
    }
    httpClient.open();
    TProtocol prot = new TBinaryProtocol(httpClient);
    THBaseService.Client client = new THBaseService.Client(prot);
    return new Pair<>(client, httpClient);
  } catch (TTransportException e) {
    throw new IOException(e);
  }

}
 

public void connect() throws Exception {
    TSocket socket = new TSocket(server.getHost(), server.getPort(), config.getConnectionTimeout());
    socket.getSocket().setTcpNoDelay(true);
    socket.getSocket().setKeepAlive(true);
    socket.getSocket().setSoLinger(true, 0);
    transport = new TFramedTransport(socket);
    transport.open();

    TProtocol protocol = new TBinaryProtocol(transport);

    client = new SuroServer.Client(protocol);
    ServiceStatus status = client.getStatus();
    if (status != ServiceStatus.ALIVE) {
        transport.close();
        throw new RuntimeException(server + " IS NOT ALIVE!!!");
    }
}
 

private boolean doCreateClient(DeviceId deviceId, String thriftServerIp, Integer thriftServerPort) {
    SafeThriftClient.Options options = new SafeThriftClient.Options(numConnectionRetries, timeBetweenRetries);

    try {
        // Make the expensive call
        TTransport transport = new TSocket(thriftServerIp, thriftServerPort);

        TProtocol protocol = new TBinaryProtocol(transport);
        // Create a client for simple_pre service.
        SimplePreLAG.Client simplePreClient = new SimplePreLAG.Client(
                new TMultiplexedProtocol(protocol, THRIFT_SERVICE_NAME));

        SimplePreLAG.Iface safeSimplePreClient = SafeThriftClient.wrap(simplePreClient,
                                                                       SimplePreLAG.Iface.class,
                                                                       options);

        Bmv2DeviceThriftClient client = new Bmv2DeviceThriftClient(deviceId, safeSimplePreClient);
        clients.put(deviceId, Pair.of(transport, client));
        return true;

    } catch (RuntimeException e) {
        log.warn("Failed to create Thrift client for BMv2 device. deviceId={}, cause={}", deviceId, e);
        return false;
    }
}
 

/**
 * To consume a list of elements
 * @param c the class of the list content
 * @param consumer the consumer that will receive the list
 * @param <T> the type of the list content
 * @return a ListConsumer that can be passed to the DelegatingFieldConsumer
 */
public static <T extends TBase<T,? extends TFieldIdEnum>> ListConsumer listOf(Class<T> c, final Consumer<List<T>> consumer) {
  class ListConsumer implements Consumer<T> {
    List<T> list;
    @Override
    public void consume(T t) {
      list.add(t);
    }
  }
  final ListConsumer co = new ListConsumer();
  return new DelegatingListElementsConsumer(struct(c, co)) {
    @Override
    public void consumeList(TProtocol protocol,
        EventBasedThriftReader reader, TList tList) throws TException {
      co.list = new ArrayList<T>();
      super.consumeList(protocol, reader, tList);
      consumer.consume(co.list);
    }
  };
}
 

private boolean readOneSet(TProtocol in, List<Action> buffer, SetType expectedType) throws TException {
  final TSet set = in.readSetBegin();
  buffer.add(new Action() {
    @Override
    public void write(TProtocol out) throws TException {
      out.writeSetBegin(set);
    }

    @Override
    public String toDebugString() {
      return "<e=" + set.elemType + ", s=" + set.size + ">{*";
    }
  });

  boolean hasFieldsIgnored = readCollectionElements(in, set.size, set.elemType, buffer, expectedType.getValues().getType());
  in.readSetEnd();
  buffer.add(SET_END);
  return hasFieldsIgnored;
}
 

private Converter newConverter(List<TProtocol> events, Type type, ThriftField field) {
  switch (field.getType().getType()) {
  case LIST:
    return new ListConverter(events, type.asGroupType(), field);
  case SET:
    return new SetConverter(events, type.asGroupType(), field);
  case MAP:
    return new MapConverter(events, type.asGroupType(), field);
  case STRUCT:
    return new StructConverter(events, type.asGroupType(), field);
  case STRING:
    return new FieldStringConverter(events, field);
  case ENUM:
    return new FieldEnumConverter(events, field);
  default:
    return new FieldPrimitiveConverter(events, field);
  }
}
 

/**
 * Verify that thrift2 client calling thrift server can get the thrift server type correctly.
 */
@Test
public void testGetThriftServerOneType() throws Exception {

  // start a thrift server
  HBaseThriftTestingUtility THRIFT_TEST_UTIL = new HBaseThriftTestingUtility();

  LOG.info("Starting HBase Thrift server One");
  THRIFT_TEST_UTIL.startThriftServer(UTIL.getConfiguration(), ThriftServerType.ONE);
  try (TTransport transport = new TSocket(InetAddress.getLocalHost().getHostName(),
      THRIFT_TEST_UTIL.getServerPort())){
    TProtocol protocol = new TBinaryProtocol(transport);
    // This is our thrift2 client.
    THBaseService.Iface client = new THBaseService.Client(protocol);
    // open the transport
    transport.open();
    assertEquals(TThriftServerType.ONE.name(), client.getThriftServerType().name());
  } finally {
    THRIFT_TEST_UTIL.stopThriftServer();
  }
}
 

@Override
protected void messageReceived(ChannelHandlerContext ctx, ThriftMessage message) throws Exception {
	ByteBuf buffer = message.getContent();
	logger.debug("msg.content:: {}", buffer);
	try {
		TNettyTransport transport = new TNettyTransport(ctx.channel(), buffer);
		TProtocolFactory protocolFactory = message.getProtocolFactory();
		TProtocol protocol = protocolFactory.getProtocol(transport);
		serverDef.nettyProcessor.process(ctx, protocol, protocol,
				new DefaultWriterListener(message, transport, ctx, serverDef));
	} catch (Throwable ex) {
		int refCount = buffer.refCnt();
		if (refCount > 0) {
			buffer.release(refCount);
		}
		throw ex;
	}
}
 

public ElementConverter(String listName, List<TProtocol> listEvents,
                        GroupType repeatedType, ThriftField thriftElement) {
  this.listEvents = listEvents;
  this.elementEvents = new ArrayList<TProtocol>();
  Type elementType = repeatedType.getType(0);
  if (elementType.isRepetition(Type.Repetition.OPTIONAL)) {
    if (ignoreNullElements) {
      LOG.warn("List " + listName +
          " has optional elements: null elements are ignored.");
    } else {
      throw new ParquetDecodingException("Cannot read list " + listName +
          " with optional elements: set " + IGNORE_NULL_LIST_ELEMENTS +
          " to ignore nulls.");
    }
  }
  elementConverter = newConverter(elementEvents, elementType, thriftElement);
}
 

@Override
public final boolean process(final TProtocol in, final TProtocol out)
    throws TException {
  final TMessage msg = in.readMessageBegin();
  final ProcessFunction<LocatorServiceImpl, ?> fn = this.fnMap
      .get(msg.name);
  if (fn != null) {
    fn.process(msg.seqid, in, out, this.inst);
    // terminate connection on receiving closeConnection
    // direct class comparison should be the fastest way
    return fn.getClass() != LocatorService.Processor.closeConnection.class;
  }
  else {
    TProtocolUtil.skip(in, TType.STRUCT);
    in.readMessageEnd();
    TApplicationException x = new TApplicationException(
        TApplicationException.UNKNOWN_METHOD, "Invalid method name: '"
            + msg.name + "'");
    out.writeMessageBegin(new TMessage(msg.name, TMessageType.EXCEPTION,
        msg.seqid));
    x.write(out);
    out.writeMessageEnd();
    out.getTransport().flush();
    return true;
  }
}
 

public static void startClient2(String ip ,int port ,int timeout) throws Exception
{
	TTransport transport = new TFramedTransport(new TSocket(ip,port,timeout));
	TProtocol protocol = new TBinaryProtocol(transport);
	leafrpc.Client client = new leafrpc.Client(protocol);
	transport.open();

	for(int i = 0; i< 1000000; i++)
	{
		client.getID("");
		if (i % 100000 == 0)
		{
			System.out.println(Thread.currentThread().getName() + " " + client.getID(""));
		}
		//ai.incrementAndGet();
	}
	transport.close();
}
 

public static Greeter.Client[] createClients(String host, int n, BenchmarkMessage msg) throws TException {
    TTransport[] transport = new TTransport[n];
    Greeter.Client[] clients = new Greeter.Client[n];

    //warmup
    for (int i = 0; i < n; i++) {
        transport[i] = new TFramedTransport(new TSocket(host, 8972));
        transport[i].open();

        TProtocol protocol = new TBinaryProtocol(transport[i]);
        clients[i] =  new Greeter.Client(protocol);
        clients[i].say(msg);
    }

    return clients;
}
 

protected void encodeRequest(Channel channel, ChannelBuffer buffer, Request request)
    throws IOException {
    Invocation invocation = (Invocation) request.getData();
    TProtocol protocol = newProtocol(channel.getUrl(), buffer);
    try {
        protocol.writeMessageBegin(new TMessage(
            invocation.getMethodName(), TMessageType.CALL, 
            thriftSeq.getAndIncrement()));
        protocol.writeStructBegin(new TStruct(invocation.getMethodName() + "_args"));
        for(int i = 0; i < invocation.getParameterTypes().length; i++) {
            Class<?> type = invocation.getParameterTypes()[i];

        }
    } catch (TException e) {
        throw new IOException(e.getMessage(), e);
    }

}
 

public static void main(String[] args) {
    TTransport trans = null;
    try {
        System.out.println("Radio Observation Disk Serializer " + radio_observationConstants.Version);
        trans = new TSimpleFileTransport("data.z", true, true);
        trans = new TZlibTransport(trans);
        trans.open();
        TProtocol proto = new TBinaryProtocol(trans);
        if (args.length > 0 && 0 == args[0].compareTo("write")) {
            WriteRadioObservation(proto, args.length > 1);
        } else if (args.length>0 && 0==args[0].compareToIgnoreCase("read")) {
            ReadRadioObservation(proto);
        } else {
            System.out.println("Usage: DiskSer (read | write [bmp])");
        }
    } catch (TException | IOException ex) {
        System.out.println("Error: " + ex.getMessage());
    }
    if (null != trans) {
        trans.close();
    }
}
 

@Nullable
private static TApplicationException readApplicationException(int seqId, ThriftFunction func,
                                                              TProtocol inputProtocol,
                                                              TMessage msg) throws TException {
    if (msg.seqid != seqId) {
        throw new TApplicationException(TApplicationException.BAD_SEQUENCE_ID);
    }

    if (!func.name().equals(msg.name)) {
        return new TApplicationException(TApplicationException.WRONG_METHOD_NAME, msg.name);
    }

    if (msg.type == TMessageType.EXCEPTION) {
        final TApplicationException appEx = TApplicationExceptions.read(inputProtocol);
        inputProtocol.readMessageEnd();
        return appEx;
    }

    return null;
}
 

protected ClientProcessData(GfxdTSocket socket, int connectionNumber,
    TProcessor proc, TTransport in, TTransport out, TProtocol inp,
    TProtocol outp, TServerEventHandler eventHandler) {
  this.clientSocket = socket;
  this.connectionNumber = connectionNumber;
  this.processor = proc;
  this.inputTransport = in;
  this.outputTransport = out;
  this.inputProtocol = inp;
  this.outputProtocol = outp;
  this.eventHandler = eventHandler;
  if (eventHandler != null) {
    this.connectionContext = eventHandler.createContext(inp, outp);
  }
  else {
    this.connectionContext = null;
  }
  this.idle = true;
}
 

public static void WriteRadioObservation(TProtocol proto,
                                         boolean writeBMP)
        throws TException, IOException {
    System.out.println("\nWriting Observations");
    System.out.println("-------------------------");
    RadioObservation ro = new RadioObservation();
    FakeInit(ro);
    if (writeBMP){
        ro.setSky_bmp(Files.readAllBytes(Paths.get("quasar.bmp")));
    }
    ro.write(proto);
    proto.getTransport().flush();
    DumpObservation(ro);
}
 

protected ChannelBuffer tryDecodeUnframedMessage(ChannelHandlerContext ctx, Channel channel, ChannelBuffer buffer,
		TProtocolFactory inputProtocolFactory) throws TException {

	int messageLength = 0;
	int messageStartReaderIndex = buffer.readerIndex();

	try {
		TNettyTransport decodeAttemptTransport = new TNettyTransport(channel, buffer, TNettyTransportType.UNFRAMED);
		int initialReadBytes = decodeAttemptTransport.getReadByteCount();
		TProtocol inputProtocol = inputProtocolFactory.getProtocol(decodeAttemptTransport);
		inputProtocol.readMessageBegin();
		TProtocolUtil.skip(inputProtocol, TType.STRUCT);
		inputProtocol.readMessageEnd();
		messageLength = decodeAttemptTransport.getReadByteCount() - initialReadBytes;
	} catch (TTransportException | IndexOutOfBoundsException e) {
		return null;
	} finally {
		if (buffer.readerIndex() - messageStartReaderIndex > maxFrameSize)
			Channels.fireExceptionCaught(ctx,
					new TooLongFrameException(String.format("Maximum frame size of %d exceeded .", maxFrameSize)));
		buffer.readerIndex(messageStartReaderIndex);
	}

	if (messageLength <= 0)
		return null;

	ChannelBuffer messageBuffer = extractFrame(buffer, messageStartReaderIndex, messageLength);
	buffer.readerIndex(messageStartReaderIndex + messageLength);
	return messageBuffer;
}
 

private void readOneSet(TProtocol in, TProtocol out) throws TException {
  final TSet set = in.readSetBegin();
  out.writeSetBegin(set);
  readCollectionElements(in, out, set.size, set.elemType);
  in.readSetEnd();
  out.writeSetEnd();
}
 

@SuppressWarnings({ "rawtypes" })
private void writeException(final TProtocol out, final TMessage msg, final WriterHandler onComplete,
		final TApplicationException x, TBase args) {
	Throwable cause = null;
	try {
		onComplete.beforeWrite(msg, args, null);
		out.writeMessageBegin(new TMessage(msg.name, TMessageType.EXCEPTION, msg.seqid));
		x.write(out);
		out.writeMessageEnd();
		out.getTransport().flush();
	} catch (Throwable e) {
		cause = e;
	}
	onComplete.afterWrite(msg, cause, TMessageType.EXCEPTION, args, null);
}
 

private void writeParquetFile(List<TBase> recordsToWrite, Configuration conf, Path parquetFile) throws IOException, InterruptedException, org.apache.thrift.TException {
  //create a test file
  final TProtocolFactory protocolFactory = new TCompactProtocol.Factory();
  final TaskAttemptID taskId = new TaskAttemptID("local", 0, true, 0, 0);
  Class writeClass = recordsToWrite.get(0).getClass();
  final ThriftToParquetFileWriter w = new ThriftToParquetFileWriter(parquetFile, ContextUtil.newTaskAttemptContext(conf, taskId), protocolFactory, writeClass);
  final ByteArrayOutputStream baos = new ByteArrayOutputStream();
  final TProtocol protocol = protocolFactory.getProtocol(new TIOStreamTransport(baos));
  for (TBase recordToWrite : recordsToWrite) {
    recordToWrite.write(protocol);
  }
  w.write(new BytesWritable(baos.toByteArray()));
  w.close();
}
 

@Test
public void voidMethodWithoutOneway() throws Exception {
    NotifyCheck notifyCheck = new NotifyCheck(100);
    final TestService.Processor<Iface> processor = new TestService.Processor<Iface>(notifyCheck);
    final TTransport transport = new FlushableMemoryBuffer(4096) {
        boolean flushed = false;

        @Override
        public void flush() throws TTransportException {
            if (!flushed) {
                flushed = true;
                try {
                    processor.process(new TBinaryProtocol(this), new TBinaryProtocol(this));
                } catch (TException e) {
                    throw new RuntimeException(e);
                }
            }
        }
    };
    com.meituan.firefly.testfirefly.TestService testService = thrift.create(com.meituan.firefly.testfirefly.TestService.class, new Thrift.SimpleTProtocolFactory() {
        @Override
        public TProtocol get() {
            return new TBinaryProtocol(transport);
        }
    });
    testService.notifyWithoutOneway(100);
    Assert.assertTrue(notifyCheck.notified);
}
 

protected ClientProcessData newClientProcessData(GfxdTSocket client) {
  TProcessor processor = processorFactory_.getProcessor(client);
  TProtocol inputProtocol = inputProtocolFactory_.getProtocol(client);
  TProtocol outputProtocol = outputProtocolFactory_.getProtocol(client);

  final int connectionNumber = connectionCounter.incrementAndGet();
  // register with ConnectionListener
  final ConnectionListener listener = connListener;
  if (listener != null) {
    listener.connectionOpened(client.getSocket(), connectionNumber);
  }

  return new ClientProcessData(client, connectionNumber, processor, client,
      client, inputProtocol, outputProtocol, getEventHandler());
}
 

private String getInetAddress(TProtocol in) {
	TTransport transport = in.getTransport();
	if (transport != null && transport instanceof TSocket) {
		Socket socket = ((TSocket) in.getTransport()).getSocket();
		return socket.getInetAddress().getHostAddress().replace('.', ':');
	} else {
		return UN_KNOWN_IP;
	}
}
 

@Override
public boolean process(TProtocol in, TProtocol out) throws TException {
	try {
		return super.process(in, out);
	} catch (TTransportException e) {
		return false; // 如果连接中断就停止服务但不抛出异常
	}
}
 

@Override
public Map<?, ?> read(TProtocol protocol) throws TException {
    TMap tmap = protocol.readMapBegin();
    HashMap hashMap = new HashMap(tmap.size);
    for (int i = 0, n = tmap.size; i < n; i++) {
        Object key = keyTypeAdapter.read(protocol);
        Object value = valueTypeAdapter.read(protocol);
        hashMap.put(key, value);
    }
    protocol.readMapEnd();
    return hashMap;
}
 

public void write(final TBase<?, ?> base) throws TException {
    final TBaseStreamNode node = new TBaseStreamNode(transport);
    node.setClassName(base.getClass().getName());
    node.setBeginPosition(transport.getBufferPosition());

    final TProtocol protocol = protocolFactory.getProtocol(transport);
    base.write(protocol);

    node.setEndPosition(transport.getBufferPosition());
    nodes.add(node);
}
 

private void readOneStruct(TProtocol in, TProtocol out) throws TException {
  final TStruct struct = in.readStructBegin();
  out.writeStructBegin(struct);
  TField field;
  while ((field = in.readFieldBegin()).type != TType.STOP) {
    out.writeFieldBegin(field);
    readOneValue(in, out, field.type);
    in.readFieldEnd();
    out.writeFieldEnd();
  }
  out.writeFieldStop();
  in.readStructEnd();
  out.writeStructEnd();
}