/*
\ * Copyright 2015 NAVER Corp.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package com.navercorp.pinpoint.plugin.thrift;
import java.nio.ByteBuffer;
import java.nio.CharBuffer;
import java.nio.charset.Charset;
import java.util.EnumMap;
import com.navercorp.pinpoint.common.Charsets;
import org.apache.thrift.TException;
import org.apache.thrift.protocol.TField;
import org.apache.thrift.protocol.TProtocol;
import org.apache.thrift.protocol.TProtocolException;
import org.apache.thrift.protocol.TType;
/**
* @author HyunGil Jeong
*/
public class ThriftRequestProperty {
private EnumMap<ThriftHeader, Object> thriftHeaders = new EnumMap<ThriftHeader, Object>(ThriftHeader.class);
// TRACE_ID
public String getTraceId() {
return this.getTraceId(null);
}
public String getTraceId(String defaultValue) {
if (this.thriftHeaders.containsKey(ThriftHeader.THRIFT_TRACE_ID)) {
return (String)this.thriftHeaders.get(ThriftHeader.THRIFT_TRACE_ID);
}
return defaultValue;
}
public void setTraceId(String traceId) {
this.thriftHeaders.put(ThriftHeader.THRIFT_TRACE_ID, traceId);
}
// SPAN_ID
public Long getSpanId() {
return this.getSpanId(null);
}
public Long getSpanId(Long defaultValue) {
if (this.thriftHeaders.containsKey(ThriftHeader.THRIFT_SPAN_ID)) {
return (Long)this.thriftHeaders.get(ThriftHeader.THRIFT_SPAN_ID);
}
return defaultValue;
}
public void setSpanId(Long spanId) {
this.thriftHeaders.put(ThriftHeader.THRIFT_SPAN_ID, spanId);
}
// PARENT_SPAN_ID
public Long getParentSpanId() {
return this.getParentSpanId(null);
}
public Long getParentSpanId(Long defaultValue) {
if (this.thriftHeaders.containsKey(ThriftHeader.THRIFT_PARENT_SPAN_ID)) {
return (Long)this.thriftHeaders.get(ThriftHeader.THRIFT_PARENT_SPAN_ID);
}
return defaultValue;
}
public void setParentSpanId(Long parentSpanId) {
this.thriftHeaders.put(ThriftHeader.THRIFT_PARENT_SPAN_ID, parentSpanId);
}
// SAMPLED
public Boolean shouldSample() {
return this.shouldSample(null);
}
public Boolean shouldSample(Boolean defaultValue) {
if (this.thriftHeaders.containsKey(ThriftHeader.THRFIT_SAMPLED)) {
return (Boolean)this.thriftHeaders.get(ThriftHeader.THRFIT_SAMPLED);
}
return defaultValue;
}
public void setShouldSample(Boolean shouldSample) {
this.thriftHeaders.put(ThriftHeader.THRFIT_SAMPLED, shouldSample);
}
// FLAGS
public Short getFlags() {
return this.getFlags(null);
}
public Short getFlags(Short defaultValue) {
if (this.thriftHeaders.containsKey(ThriftHeader.THRIFT_FLAGS)) {
return (Short)this.thriftHeaders.get(ThriftHeader.THRIFT_FLAGS);
}
return defaultValue;
}
public void setFlags(Short flags) {
this.thriftHeaders.put(ThriftHeader.THRIFT_FLAGS, flags);
}
// PARENT_APPLICATION_NAME
public String getParentApplicationName() {
return this.getParentApplicationName(null);
}
public String getParentApplicationName(String defaultValue) {
if (this.thriftHeaders.containsKey(ThriftHeader.THRIFT_PARENT_APPLICATION_NAME)) {
return (String)this.thriftHeaders.get(ThriftHeader.THRIFT_PARENT_APPLICATION_NAME);
}
return defaultValue;
}
public void setParentApplicationName(String parentApplicationName) {
this.thriftHeaders.put(ThriftHeader.THRIFT_PARENT_APPLICATION_NAME, parentApplicationName);
}
// PARENT_APPLICATION_TYPE
public Short getParentApplicationType() {
return this.getParentApplicationType(null);
}
public Short getParentApplicationType(Short defaultValue) {
if (this.thriftHeaders.containsKey(ThriftHeader.THRIFT_PARENT_APPLICATION_TYPE)) {
return (Short)this.thriftHeaders.get(ThriftHeader.THRIFT_PARENT_APPLICATION_TYPE);
}
return defaultValue;
}
public void setParentApplicationType(Short parentApplicationType) {
this.thriftHeaders.put(ThriftHeader.THRIFT_PARENT_APPLICATION_TYPE, parentApplicationType);
}
// ACCEPTOR_HOST
public String getAcceptorHost() {
return getAcceptorHost(ThriftConstants.UNKNOWN_ADDRESS);
}
public String getAcceptorHost(String acceptorHost) {
if (this.thriftHeaders.containsKey(ThriftHeader.THRIFT_HOST)) {
return (String)this.thriftHeaders.get(ThriftHeader.THRIFT_HOST);
}
return acceptorHost;
}
public void setAcceptorHost(String acceptorHost) {
this.thriftHeaders.put(ThriftHeader.THRIFT_HOST, acceptorHost);
}
public void setTraceHeader(ThriftHeader headerKey, Object value) throws TException {
byte headerType = headerKey.getType();
if (headerType == TType.STRING) {
// skipped Strings are read as byte buffer.
// see org.apache.thrift.protocol.TProtocolUtil.skip(TProtocol, byte, int)
this.thriftHeaders.put(headerKey, byteBufferToString((ByteBuffer)value));
} else if (headerType == TType.I64) {
this.thriftHeaders.put(headerKey, (Long)value);
} else if (headerType == TType.I16) {
this.thriftHeaders.put(headerKey, (Short)value);
} else if (headerType == TType.BOOL) {
this.thriftHeaders.put(headerKey, (Boolean)value);
} else {
throw new TProtocolException("Invalid pinpoint header type - " + headerType);
}
}
public void writeTraceHeader(ThriftHeader headerKey, TProtocol oprot) throws TException {
Object headerValue = this.thriftHeaders.get(headerKey);
if (headerValue == null) {
return;
}
byte headerType = headerKey.getType();
TField traceField = new TField(headerKey.name(), headerKey.getType(), headerKey.getId());
oprot.writeFieldBegin(traceField);
try {
if (headerType == TType.STRING) {
// these will be read as byte buffer although it's probably safe to just use writeString here.
// see org.apache.thrift.protocol.TProtocolUtil.skip(TProtocol, byte, int)
oprot.writeBinary(stringToByteBuffer((String)headerValue));
} else if (headerType == TType.I64) {
oprot.writeI64((Long)headerValue);
} else if (headerType == TType.I16) {
oprot.writeI16((Short)headerValue);
} else if (headerType == TType.BOOL) {
oprot.writeBool((Boolean)headerValue);
} else {
throw new TProtocolException("Invalid pinpoint header type - " + headerType);
}
} finally {
oprot.writeFieldEnd();
}
}
private static final Charset HEADER_CHARSET_ENCODING = Charsets.UTF_8;
private static ByteBuffer stringToByteBuffer(String s) {
return ByteBuffer.wrap(s.getBytes(HEADER_CHARSET_ENCODING));
}
private static String byteBufferToString(ByteBuffer buf) {
CharBuffer charBuffer = HEADER_CHARSET_ENCODING.decode(buf);
return charBuffer.toString();
}
@Override
public String toString() {
return this.thriftHeaders.toString();
}
}