org.apache.thrift.TBase Java Examples

@Override
public void append(byte[] partitionKey, TBase thriftMessage, long timeNanos) throws TException {
  LoggingAuditHeaders headers = null;
  try {
    if (enableLoggingAudit && shouldAudit()) {
      headers = auditHeadersGenerator.generateHeaders();
      if (this.loggingAuditHeadersField != null) {
        thriftMessage.setFieldValue(this.loggingAuditHeadersField, headers);
        OpenTsdbMetricConverter.incr(AUDIT_THRIFT_LOGGER_HEADERS_ADDED_TO_ORIGINAL_COUNT,
            "topic=" + topic, "host=" + HOST_NAME);
      }
    }
    byte[] messageBytes = ThriftCodec.getInstance().serialize(thriftMessage);
    append(partitionKey, messageBytes, timeNanos, headers);
  } catch (TException e) {
    OpenTsdbMetricConverter.incr(THRIFT_LOGGER_ERROR_TEXCEPTION,
        "topic=" + topic, "host=" + HOST_NAME);
    throw e;
  }
}
 

public MockResult(String methodName, TBase value) {
  this.methodName = methodName;
  this.success = value;
  this.structDesc = new TStruct(methodName+"_result");

  //init metaDatMap
  Map<_Fields, FieldMetaData> tmpMap = new EnumMap<>(_Fields.class);
  tmpMap.put(_Fields.SUCCESS,
             new FieldMetaData(SUCCESS_NAME, TFieldRequirementType.DEFAULT,
                               new FieldValueMetaData(TType.STRUCT          , value.getClass().getCanonicalName())));
  metaDataMap = Collections.unmodifiableMap(tmpMap);
  FieldMetaData.addStructMetaDataMap(MockResult.class, metaDataMap);

  schemes.put(StandardScheme.class, new MockResultStandardSchemeFactory(structDesc));
  schemes.put(TupleScheme.class, new MockResultTupleSchemeFactory());

}
 

private static void invokeAsynchronously(Object impl, ThriftFunction func, TBase<?, ?> args,
                                         CompletableRpcResponse reply) throws TException {

    final AsyncProcessFunction<Object, TBase<?, ?>, Object> f = func.asyncFunc();
    if (func.isOneWay()) {
        f.start(impl, args, ONEWAY_CALLBACK);
        reply.complete(null);
    } else {
        f.start(impl, args, new AsyncMethodCallback<Object>() {
            @Override
            public void onComplete(Object response) {
                reply.complete(response);
            }

            @Override
            public void onError(Exception e) {
                reply.completeExceptionally(e);
            }
        });
    }
}
 

/** Create request that sends no out-of-band payloads. */
public static <ThriftRequest extends TBase<?, ?>>
    HybridThriftRequestHandler<ThriftRequest> createWithoutPayloads(ThriftRequest request) {
  return new HybridThriftRequestHandler<ThriftRequest>(request) {

    @Override
    public long getTotalPayloadsSizeBytes() {
      return 0;
    }

    @Override
    public int getNumberOfPayloads() {
      return 0;
    }

    @Override
    public InputStream getPayloadStream(int index) {
      throw new IllegalStateException();
    }
  };
}
 

private void addTSpan(TBase<?, ?> base) throws TException {
    final TSpan span = (TSpan) base;
    if (span.getSpanEventList() == null) {
        write(base);
        return;
    }

    try {
        for (TSpanEvent e : span.getSpanEventList()) {
            eventStream.write(e);
        }
        write(span, FIELD_NAME_SPAN_EVENT_LIST, eventStream.split(chunkSize));
        while (!eventStream.isEmpty()) {
            final TSpanChunk spanChunk = toSpanChunk(span);
            write(spanChunk, FIELD_NAME_SPAN_EVENT_LIST, eventStream.split(chunkSize));
        }
    } finally {
        eventStream.clear();
    }
}
 

@Test
public void addBodyFactory() {
    TypeLocatorBuilder<TBase<?, ?>> typeLocatorBuilder = new TypeLocatorBuilder<TBase<?, ?>>();

    typeLocatorBuilder.addBodyFactory((short) 1, new BodyFactory<TBase<?, ?>>() {
        @Override
        public TBase<?, ?> getObject() {
            return new TSpan();
        }
    });

    typeLocatorBuilder.addBodyFactory((short) 3, new BodyFactory<TBase<?, ?>>() {
        @Override
        public TBase<?, ?> getObject() {
            return new TSpanEvent();
        }
    });

    TypeLocator<TBase<?, ?>> build = typeLocatorBuilder.build();
    Assert.assertNotNull(build.bodyLookup((short) 1));
    Assert.assertNull(build.bodyLookup((short) 5));

    Assert.assertEquals(build.headerLookup((short) 1).getType(), 1);
    Assert.assertEquals(build.headerLookup((short) 3).getType(), 3);
}
 

public NioUDPDataSender(String host, int port, String threadName, int queueSize, int timeout, int sendBufferSize,
                        MessageConverter<TBase<?, ?>> messageConverter) {
    Assert.requireNonNull(host, "host");
    Assert.requireNonNull(threadName, "threadName");
    Assert.isTrue(queueSize > 0, "queueSize");
    Assert.isTrue(timeout > 0, "timeout");
    Assert.isTrue(sendBufferSize > 0, "sendBufferSize");

    this.messageConverter = Assert.requireNonNull(messageConverter, "messageConverter");

    // TODO If fail to create socket, stop agent start
    logger.info("NioUDPDataSender initialized. host={}, port={}", host, port);
    this.datagramChannel = createChannel(host, port, timeout, sendBufferSize);

    HeaderTBaseSerializerFactory2 serializerFactory = new HeaderTBaseSerializerFactory2();
    this.serializer = serializerFactory.createSerializer();

    ByteBufferFactory bufferFactory = ByteBufferFactoryLocator.getFactory(ByteBufferType.DIRECT);
    ByteBuffer byteBuffer = bufferFactory.getBuffer(UDP_MAX_PACKET_LENGTH);
    this.byteBufferOutputStream = new ByteBufferOutputStream(byteBuffer);

    this.executor = createAsyncQueueingExecutor(queueSize, threadName);
}
 

/**
 * Converts the specified {@code result} into a Java object.
 */
@Nullable
public Object getResult(TBase<?, ?> result) throws TException {
    for (TFieldIdEnum fieldIdEnum : exceptionFields()) {
        if (ThriftFieldAccess.isSet(result, fieldIdEnum)) {
            throw (TException) ThriftFieldAccess.get(result, fieldIdEnum);
        }
    }

    final TFieldIdEnum successField = successField();
    if (successField == null) { //void method
        return null;
    } else if (ThriftFieldAccess.isSet(result, successField)) {
        return ThriftFieldAccess.get(result, successField);
    } else {
        throw new TApplicationException(
                TApplicationException.MISSING_RESULT,
                result.getClass().getName() + '.' + successField.getFieldName());
    }
}
 

@SuppressWarnings({ "rawtypes" })
@Override
public void beforeWrite(TMessage msg, TBase args, TBase result) {
	// reuse message's buffer when write? yes, we use the pool.
	ByteBuf readedBuf = message.getContent();
	int refCount = readedBuf.refCnt();
	if (refCount > 0) {
		readedBuf.release(refCount);
	}
	// voidMethod's return message is very short
	int initialCapacity = serverDef.trafficForecast.getInitBytesForWrite(msg.name);
	// logger.debug("initialCapacity = {} , msg = {}",initialCapacity, msg);
	ByteBuf buf = ctx.alloc().buffer(initialCapacity, serverDef.maxFrameSize);
	message.setContent(buf).beforeWrite(ctx);
	transport.setOutputBuffer(buf);
}
 

public String encryptToken(TBase<?, ?> token, byte[] tokenAESKey, byte[] tokenSipHashKey) throws TException {
  // Serialize the  thrift token into byte array
  byte[] serialized = serializer.serialize(token);

  // Calculate the SIP
  long sip = SipHashInline.hash24_palindromic(tokenSipHashKey, serialized);

  //Create the token byte buffer
  ByteBuffer buffer = ByteBuffer.allocate(8 + serialized.length);
  buffer.order(ByteOrder.BIG_ENDIAN);
  // adds the sip
  buffer.putLong(sip);
  // adds the thrift token
  buffer.put(serialized);

  // Wrap the TOKEN
  byte[] wrappedData = CryptoUtils.wrap(tokenAESKey, buffer.array());

  String accessToken = new String(OrderPreservingBase64.encode(wrappedData));

  return accessToken;
}
 

@Override
public String dump(LogMessage logMessage, boolean noTimestamp, String thriftSchema)
    throws Exception {
  JsonDumperMessage log = new JsonDumperMessage();

  if (!noTimestamp) {
    log.timeStampInNanos = logMessage.getTimestampInNanos();
  }
  if (thriftSchema != null && !thriftSchema.isEmpty()) {
    TBase tObj = deserializeBytesToThriftObj(logMessage.getMessage(), thriftSchema);
    log.deserializedMessage = tObj;
  } else {
    log.message = new String(logMessage.getMessage());
  }
  return gson.toJson(log);
}
 

/**
 * <ul> steps:
 * <li>Writes using the thrift mapping
 * <li>Reads using the pig mapping
 * <li>Use Elephant bird to convert from thrift to pig
 * <li>Check that both transformations give the same result
 * @param o the object to convert
 * @throws TException
 */
public static <T extends TBase<?,?>> void validateSameTupleAsEB(T o) throws TException {
  final ThriftSchemaConverter thriftSchemaConverter = new ThriftSchemaConverter();
  @SuppressWarnings("unchecked")
  final Class<T> class1 = (Class<T>) o.getClass();
  final MessageType schema = thriftSchemaConverter.convert(class1);

  final StructType structType = ThriftSchemaConverter.toStructType(class1);
  final ThriftToPig<T> thriftToPig = new ThriftToPig<T>(class1);
  final Schema pigSchema = thriftToPig.toSchema();
  final TupleRecordMaterializer tupleRecordConverter = new TupleRecordMaterializer(schema, pigSchema, true);
  RecordConsumer recordConsumer = new ConverterConsumer(tupleRecordConverter.getRootConverter(), schema);
  final MessageColumnIO columnIO = new ColumnIOFactory().getColumnIO(schema);
  ParquetWriteProtocol p = new ParquetWriteProtocol(new RecordConsumerLoggingWrapper(recordConsumer), columnIO, structType);
  o.write(p);
  final Tuple t = tupleRecordConverter.getCurrentRecord();
  final Tuple expected = thriftToPig.getPigTuple(o);
  assertEquals(expected.toString(), t.toString());
  final MessageType filtered = new PigSchemaConverter().filter(schema, pigSchema);
  assertEquals(schema.toString(), filtered.toString());
}
 

/**
 * Messages to be consumed by Secor must have an i64 timestamp
 * as the first field.
 */
private boolean validateThriftMessageForSecor(TBase thriftMessage) {
  TFieldIdEnum fieldEnum = thriftMessage.fieldForId(1);
  if (fieldEnum == null) {
    return false;
  }
  if (!thriftMessage.isSet(fieldEnum)) {
    return false;
  }
  Object value = thriftMessage.getFieldValue(fieldEnum);
  if (value == null || value.getClass() != Long.class) {
    return false;
  }
  return true;
}
 

private static String thriftToString(TBase object)
{
    assert object != null;
    // this is so awful it's kind of cool!
    TSerializer serializer = new TSerializer(new TBinaryProtocol.Factory());
    try
    {
        return Hex.bytesToHex(serializer.serialize(object));
    }
    catch (TException e)
    {
        throw new RuntimeException(e);
    }
}
 

public void isSupportTest_Inheritance() throws TException {
    TypeLocator<TBase<?, ?>> registry = TCommandRegistry.build(TCommandTypeVersion.V_1_0_2_SNAPSHOT);

    boolean isSupport = registry.isSupport(TResultEx.class);
    Assert.assertTrue(isSupport);

    isSupport = registry.isSupport(TCommandTransferResponse.class);
    Assert.assertFalse(isSupport);
}
 

/**
 * Sets the exception field of the specified {@code result} to the specified {@code cause}.
 */
public boolean setException(TBase<?, ?> result, Throwable cause) {
    final Class<?> causeType = cause.getClass();
    for (Entry<Class<Throwable>, TFieldIdEnum> e : exceptionFields.entrySet()) {
        if (e.getKey().isAssignableFrom(causeType)) {
            ThriftFieldAccess.set(result, e.getValue(), cause);
            return true;
        }
    }
    return false;
}
 

public static <T extends TBase> String ThriftToJSON(T thrift) {
  TSerializer serializer = new TSerializer(new TSimpleJSONProtocol.Factory());
  try {
    return serializer.toString(thrift);
  } catch (TException e) {
  }
  throw new IllegalStateException("Convert to json failed : " + thrift);
}
 

@Override
public byte[] serialize(Object object) {
    try {
        return Utils.gzip(new TSerializer().serialize((TBase) object));
    } catch (TException e) {
        throw new RuntimeException(e);
    }
}
 

public void serialize(TBase<?, ?> base, OutputStream outputStream) throws TException {
    tOutputStreamTransport.open(outputStream);
    try {
        final Header header = tBaseLocator.headerLookup(base);
        if (header == null) {
            throw new TException("header must not be null base:" + base);
        }
        HeaderUtils.writeHeader(protocol, header);
        base.write(protocol);
    } finally {
        tOutputStreamTransport.close();
    }
}
 

public static void shouldGetProjectedSchema(String deprecatedFilterDesc, String strictFilterDesc, String expectedSchemaStr, Class<? extends TBase<?,?>> thriftClass) {
  MessageType depRequestedSchema = getDeprecatedFilteredSchema(deprecatedFilterDesc, thriftClass);
  MessageType strictRequestedSchema = getStrictFilteredSchema(strictFilterDesc, thriftClass);
  MessageType expectedSchema = parseMessageType(expectedSchemaStr);
  assertEquals(expectedSchema, depRequestedSchema);
  assertEquals(expectedSchema, strictRequestedSchema);
}
 

@Override
public TBase<?, ?> deserializeResponse(byte[] objectData, Message<TBase<?, ?>> defaultValue) {
    Message<TBase<?, ?>> message = SerializationUtils.deserialize(objectData, commandDeserializerFactory, defaultValue);
    if (message == null) {
        return null;
    }
    return message.getData();
}
 

public boolean deserialize(@SuppressWarnings("rawtypes") TBase obj, byte[] payload) {
  try {
    _deserializer.deserialize(obj, payload);
  } catch (TException e) {
    LOGGER.error("Unable to deserialize to object :" + obj, e);
    return false;
  }
  return true;
}
 

@Override
public String guessServiceMethodName(Object exampleRequest) {
    final TBase<?, ?> exampleTBase = asTBase(exampleRequest);
    if (exampleTBase == null) {
        return null;
    }

    final String typeName = exampleTBase.getClass().getName();
    return typeName.substring(typeName.lastIndexOf('$') + 1,
                              typeName.length() - REQUEST_STRUCT_SUFFIX.length());
}
 

/**
 * Deserialize the Thrift object from a byte array.
 *
 * @param base The object to read into
 * @param bytes The array to read from
 */
public void deserialize(TBase base, byte[] bytes) throws TException {
  try {
    if (bytes.length == 0) {
      return;
    }
    if (bytes[0] == PrefixedSerializer.SECRET_BYTE) {
      if (bytes.length == 1) {
        throw new TException("Unknown prefixed protocol with byte length 1.");
      }
      switch (bytes[1]) {
        case PrefixedSerializer.COMPACT_PROTOCOL_BYTE:
          transport.reset(bytes, 2, bytes.length - 2);
          base.read(protocol);
          break;
        default:
          throw new TException("Unknown protocol with byte: " + bytes[1]);
      }
    } else {
      // Default to TBinaryProtocol decoder.
      getInstance().decoder.get().deserialize(base, bytes);
    }
  } finally {
    transport.reset(null, 0, 0);
    protocol.reset();
  }
}
 

@Override
public boolean isSupportCommand(TBase command) {
    for (TCommandType supportCommand : supportCommandList) {
        if (supportCommand.getClazz() == command.getClass()) {
            return true;
        }
    }

    TCommandTypeVersion commandVersion = TCommandTypeVersion.getVersion(agentInfo.getVersion());
    if (commandVersion.isSupportCommand(command)) {
        return true;
    }

    return false;
}
 

@Override
public <R extends TBase> R getResponse(Class<R> clazz, R defaultValue) {
    if (clazz.isInstance(response)) {
        return (R) response;
    }

    return defaultValue;
}
 

public byte[] serialize(AgentEventType agentEventType, Object eventMessage) throws UnsupportedEncodingException {
    if (agentEventType == null) {
        throw new NullPointerException("agentEventType");
    }

    Class<?> eventMessageType = agentEventType.getMessageType();
    if (eventMessageType == Void.class) {
        return EMPTY_BYTES;
    } else {
        if (eventMessage == null) {
            throw new NullPointerException("eventMessage of type [" + eventMessageType.getName() + "] expected, but was null");
        }
    }

    if (!eventMessageType.isAssignableFrom(eventMessage.getClass())) {
        throw new IllegalArgumentException("Unexpected eventMessage of type [" + eventMessage.getClass().getName() + "] received. Expected : ["
                + eventMessageType.getClass().getName() + "]");
    }

    if (eventMessage instanceof TBase) {
        for (SerializerFactory serializerFactory : serializerFactoryList) {
            if (serializerFactory.isSupport(eventMessage)) {
                try {
                    return SerializationUtils.serialize((TBase<?, ?>) eventMessage, serializerFactory);
                } catch (TException e) {
                    throw new UnsupportedEncodingException(e.getMessage());
                }
            }
        }
    } else if (eventMessage instanceof String) {
        return BytesUtils.toBytes((String) eventMessage);
    }
    throw new UnsupportedEncodingException("Unsupported event message type [" + eventMessage.getClass().getName() + "]");
}
 

@Test
@SuppressWarnings({ "rawtypes", "unchecked" })
public void handler_should_ignore_request_events_with_unsupported_message_types() throws Exception {
    // given
    final TCommandEcho mismatchingResponse = new TCommandEcho();
    final byte[] mismatchingResponseBody = new byte[0];

    final TCommandTransfer tCommandTransfer = new TCommandTransfer();
    tCommandTransfer.setAgentId(TEST_AGENT_ID);
    tCommandTransfer.setStartTime(TEST_START_TIMESTAMP);

    final TCommandTransferResponse tCommandTransferResponse = new TCommandTransferResponse();
    tCommandTransferResponse.setRouteResult(TRouteResult.OK);
    tCommandTransferResponse.setPayload(mismatchingResponseBody);

    final ResponseEvent responseEvent = new ResponseEvent(tCommandTransfer, null, 0, tCommandTransferResponse);

    ArgumentCaptor<AgentEventBo> argCaptor = ArgumentCaptor.forClass(AgentEventBo.class);
    HeaderTBaseDeserializer deserializer = mock(HeaderTBaseDeserializer.class);
    when(this.deserializerFactory.createDeserializer()).thenReturn(deserializer);
    Message<TBase<?, ?>> message = new DefaultMessage<>(new HeaderV1((short)1000), HeaderEntity.EMPTY_HEADER_ENTITY, mismatchingResponse);
    when(deserializer.deserialize(mismatchingResponseBody)).thenReturn(message);
    // when
    this.agentEventHandlingFilter.handleResponseEvent(responseEvent, TEST_EVENT_TIMESTAMP);
    // then
    verify(this.agentEventDao, never()).insert(argCaptor.capture());
}
 

@Override
public ResultResponse toMessage(Object object) {
    if (object instanceof ResponseMessage) {
        final ResponseMessage responseMessage = (ResponseMessage) object;
        final byte[] byteMessage = responseMessage.getMessage();
        final Message<TBase<?, ?>> message = SerializationUtils.deserialize(byteMessage, HeaderTBaseDeserializerFactory.DEFAULT_FACTORY, null);
        if (message == null) {
            throw new IllegalArgumentException("message is null. response message=" + responseMessage);
        }

        final TBase<?, ?> tbase = message.getData();
        if (!(tbase instanceof TResult)) {
            throw new IllegalArgumentException("invalid message data. response message=" + responseMessage + ", data=" + tbase.getClass());
        }

        final TResult result = (TResult) tbase;
        return new ResultResponse() {
            @Override
            public boolean isSuccess() {
                return result.isSuccess();
            }

            @Override
            public String getMessage() {
                return result.getMessage();
            }
        };
    }
    return null;
}
 

private ServerRequest<TBase<?, ?>> newServerRequest(Message<TBase<?, ?>> message, InetSocketAddress remoteSocketAddress) {
    final String remoteAddress = remoteSocketAddress.getAddress().getHostAddress();
    final int remotePort = remoteSocketAddress.getPort();

    ServerRequest<TBase<?, ?>> tBaseDefaultServerRequest = new DefaultServerRequest<>(message, remoteAddress, remotePort);
    return tBaseDefaultServerRequest;
}