全链路追踪关键技术-Thrift协议Header添加

首先对于远程通信的链路埋点来说，有两个关键信息需要传递，即 traceId 与 spanId(parent spanId)。

通信方式目前链路跟踪框架只支持 Http 一种。而 Http 利用 header 来传递 trace 信息也十分方便(traceId，spanId)。Client 端采用封装过的 httpClient 或者在 soa 中埋点，均往 header 中塞入了 trace 信息。Server 端采用 spring-web-servlet-filter，本质也就是个拦截器来读取 request 中的 header 信息。

而对于 Thrift 来说，并没有 header 这样的结构用于传递，所以需要另行开发。

Thrift-java 的通信流程

所有 Rpc 的流程都差不多，绿色部分表示可以扩展的点。

而对于实现调用链埋点，也就是 trace 信息传递这一简单的功能来说，可以通过扩展 Tprotocol 和 TProcessor，在通信的时候多带一个 trace 信息进去。

协议改造

改造过程第一个部分是对 thrift 本身的协议做一些改造。

Thrift 消息格式

观察 thrift 源码可知消息格式主要分为四种：

Message，读写消息头，消息头包含了版本号(version)，方法名(name)，序列号(seqid)等信息
Struct，将 RPC 方法的参数/返回值封装成结构体，读写结构体即表示要读写 RPC 方法参数了
Field，每一个参数都被抽象成 Field，Field 主要包含了字段的索引信息，类型信息等
Type，即读写各种具体的数据

预研方案

因为我们要扩展 TProtocol，而目前业务在使用的是 TBinaryProtocol，这种 protocol 的 struct 是空的，也就是 message 后面直接是 field。生成的消息大致如下：

那么我们可以选一个位置塞入我们的 trace 信息，塞哪儿呢，也就是以下各种方案的不同之处。

方案 1 & 方案 2
这两个方案差不多所以放一起。

其实 Thrift 本身有对 Protocol/Processor 做过一些扩展，参考 TMultiplexedProtocol 以及 TMultiplexedProcessor，可以略窥一二。
TMultiplexedProcessor:


public boolean process(TProtocol iprot, TProtocol oprot) throws TException {
    TMessage message = iprot.readMessageBegin();
    if (message.type != 1 && message.type != 4) {
        throw new TException("This should not have happened!?");
    } else {
        int index = message.name.indexOf(":");
        if (index < 0) {
            throw new TException("Service name not found in message name: " + message.name + ".  Did you " + "forget to use a TMultiplexProtocol in your client?");
        } else {
            String serviceName = message.name.substring(0, index);
            TProcessor actualProcessor = (TProcessor)this.SERVICE_PROCESSOR_MAP.get(serviceName);
            if (actualProcessor == null) {
                throw new TException("Service name not found: " + serviceName + ".  Did you forget " + "to call registerProcessor()?");
            } else {
                TMessage standardMessage = new TMessage(message.name.substring(serviceName.length() + ":".length()), message.type, message.seqid);
                return actualProcessor.process(new TMultiplexedProcessor.StoredMessageProtocol(iprot, standardMessage), oprot);
            }
        }
    }
}

TMultiplexedProtocol:

public void writeMessageBegin(TMessage tMessage) throws TException {
    if (tMessage.type != 1 && tMessage.type != 4) {
        super.writeMessageBegin(tMessage);
    } else {
        super.writeMessageBegin(new TMessage(this.SERVICE_NAME + ":" + tMessage.name, tMessage.type, tMessage.seqid));
    }

}

很显然，它对 message 里面的 name 字段做了扩展。那么参考它的做法我们这儿也将 trace 信息扩展到 name 字段中去，trace 信息用啥样的格式区分于一二方案：

{httpHeader 文本格式，traceId 和 spanId 分别作为 header 属性塞入}:原 name 信息
traceId:spandId:原 name 信息

以上两种方式优缺点很明显，也不用多说了。

方案 3
考虑到 thrift 读取 field 的一些特性：

Thrift 生成的代码在读写字节流时，都是按照生成的 TField 的索引号去判断，然后读取的
Thrift 提供了 skip 和 stop 解析 Filed 的机制

那么其实可以对 field 做一个扩展：

我们自定义一个 attachment，支持 kv 的 map 格式（当然你要像方案 2 一样只加 traceId 和 spanId 也可以），然后把 trace 信息传进去。

然后我们去实现读写 filed0 的方法，以及将字节流复位的 reset 方法就好了(复位后不影响 thrift 的正常执行流程)。

实际改造方案

考虑到 1，2 方案在实际使用过程中会自行拼装&解析处理 name 属性，在 thrift 这个环境下显得比较低效和另类，所以选择使用方案 3 进行改造，全部使用 thrift 原生的读写方法。

将 trace 信息塞入到自定义 headInfo 中，目前数据包格式如下：

埋点改造

现在我们知道怎么在数据中额外携带我们的自定义消息了(traceInfo)，但何时何地将信息埋进去是另一个问题，也就是关于埋点的改造，结合 thrift 通信流程考虑。

我们自定义了两种 protocol，即 TTraceClientProtocol 以及 TTraceServerProtocol，这两种均继承于 TBinaryProtocol，分别用于 client 与 server 端处理数据。

客户端

由于 thrift 使用插件根据 thrift 文件生成相应的 java 代码，会动态继承于 TServiceClient，并且实现业务的 Iface 接口，所以很难在 client 层面对 TServiceClient 做扩展。

所以客户端的埋点工作均完成于 TTraceClientProtocol 中。通过重写 writeMessageBegin 方法，在一开始就开启一个 trace 的生命周期，接着在 message 信息(version/name/seqId)写完之后，生成 trace 信息，将其作为 headInfo 写入到 field0 中去。

而通过重写 readMessageBegin 方法，在读取到返回信息时候，根据 Tmessage 的类型判断是正常返回还是异常信息，从而决定是否记录 trace 异常信息，最终结束 trace 生命周期。

具体实现代码如下：

public class TTraceClientProtocol extends TBinaryProtocol {

    private Map<String, String> HEAD_INFO;

    public TTraceClientProtocol(TTransport transport) {
        super(transport);
        HEAD_INFO = new HashMap<>();
    }

    //clientSend
    @Override
    public void writeMessageBegin(TMessage tMessage) throws TException {
        //trace start
        TraceUtils.startLocalTracer("rpc.thrift start");

        String methodName = tMessage.name;
        TraceUtils.submitAdditionalAnnotation(Constants.TRACE_THRIFT_METHOD, methodName);
        TTransport transport = this.getTransport();
        String hostAddress = ((TSocket) transport).getSocket().getRemoteSocketAddress().toString();
        TraceUtils.submitAdditionalAnnotation(Constants.TRACE_THRIFT_SERVER, hostAddress);

        super.writeMessageBegin(tMessage);
        //write trace header to field0
        writeFieldZero();
    }


    public void writeFieldZero() throws TException {
        TField TRACE_HEAD = new TField("traceHeader", TType.MAP, (short) 0);
        this.writeFieldBegin(TRACE_HEAD);
        {
            Map<String, String> traceInfo = genTraceInfo();
            this.writeMapBegin(new TMap(TType.STRING, TType.STRING, traceInfo.size()));
            for (Map.Entry<String, String> entry : traceInfo.entrySet()) {
                this.writeString(entry.getKey());
                this.writeString(entry.getValue());
            }
            this.writeMapEnd();
        }
        this.writeFieldEnd();
    }

    private Map<String, String> genTraceInfo() {
        //gen trace info
        ...
       return HEAD_INFO;
    }

    //clientReceive
    @Override
    public TMessage readMessageBegin() throws TException {
        TMessage tMessage = super.readMessageBegin();
        if (tMessage.type == TMessageType.EXCEPTION) {
            TApplicationException x = TApplicationException.read(this);
                TraceUtils.submitAdditionalAnnotation(Constants.TRACE_THRIFT_EXCEPTION, StringUtil.trimNewlineSymbolAndRemoveExtraSpace(x.getMessage()));
            TraceUtils.endAndSendLocalTracer();
        } else if (tMessage.type == TMessageType.REPLY) {
            TraceUtils.endAndSendLocalTracer();
        }
        return tMessage;
    }

}

服务端

服务端的埋点工作完成在 TTraceProcessor 中。

首先为了不影响字节流的正常读取需要依靠 BuffedIoInputStream 的 remark/reset 功能进行字节流的复位，显然在一开始需要先 remark 一下。

接着开始读取数据，在 message 信息读取完毕之后，使用 TTraceServerProtocol 读取 field0 的数据，也就是自定义 headInfo，此处是一些 trace 信息，那么处理完 trace 数据后开启一个 trace 的生命周期。

在正常的 processor 进行操作之前需要用 reset 功能先将字节流复位，接着转交给真正的 processor 去处理，在获取到 result 返回数据后结束 trace 的生命周期。

TTraceServerProtocol 的实现如下：

public class TTraceServerProtocol extends TBinaryProtocol {

    private Map<String, String> HEAD_INFO;

    public TTraceServerProtocol(TTransport transport) {
        super(transport);
        HEAD_INFO = new HashMap<>();
    }


    public boolean readFieldZero() throws TException {
        TField schemeField = this.readFieldBegin();

        if (schemeField.id == 0 && schemeField.type == TType.MAP) {
            TMap _map = this.readMapBegin();
            HEAD_INFO = new HashMap<>(2 * _map.size);
            for (int i = 0; i < _map.size; ++i) {
                String key = this.readString();
                String value = this.readString();
                HEAD_INFO.put(key, value);
            }
            this.readMapEnd();
        }
        this.readFieldEnd();
        return HEAD_INFO.size() > 0;
    }

    public Map<String, String> getHead() {
        return HEAD_INFO;
    }

    public void markTFramedTransport(TProtocol in) {
        try {
            Field tioInputStream = TIOStreamTransportFieldsCache.getInstance().getTIOInputStream();
            if (tioInputStream == null){
                return;
            }
            BufferedInputStream inputStream = (BufferedInputStream) tioInputStream.get(in.getTransport());
            inputStream.mark(0);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }


    /*
     * 重置TFramedTransport流，不影响Thrift原有流程
     */
    public void resetTFramedTransport(TProtocol in) {
        try {
            Field tioInputStream = TIOStreamTransportFieldsCache.getInstance().getTIOInputStream();
            if (tioInputStream == null){
                return;
            }
            BufferedInputStream inputStream = (BufferedInputStream) tioInputStream.get(in.getTransport());
            inputStream.reset();
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    private static class TIOStreamTransportFieldsCache {
        private static TIOStreamTransportFieldsCache instance;
        private final Field inputStream_;
        private final String TIOStreamTransport_inputStream_ = "inputStream_";

        private TIOStreamTransportFieldsCache() throws Exception {
            inputStream_ = TIOStreamTransport.class
                    .getDeclaredField(TIOStreamTransport_inputStream_);
            inputStream_.setAccessible(true);
        }

        public static TIOStreamTransportFieldsCache getInstance()
                throws Exception {
            if (instance == null) {
                synchronized (TIOStreamTransportFieldsCache.class) {
                    if (instance == null) {
                        instance = new TIOStreamTransportFieldsCache();
                    }
                }
            }
            return instance;
        }

        public Field getTIOInputStream() {
            return inputStream_;
        }
    }


    public static class Factory implements TProtocolFactory {

        public TProtocol getProtocol(TTransport trans) {
            return new TTraceServerProtocol(trans);
        }
    }

}

TTraceProcessor 的实现如下：

public class TTraceProcessor implements TProcessor {

    private TProcessor realProcessor;

    public TTraceProcessor(TProcessor realProcessor) {
        this.realProcessor = realProcessor;
    }

    public boolean process(TProtocol in, TProtocol out) throws TException {
        if (in instanceof TTraceServerProtocol) {
            TTraceServerProtocol serverProtocol = (TTraceServerProtocol) in;
            serverProtocol.markTFramedTransport(in);
            TMessage tMessage = serverProtocol.readMessageBegin();
            serverProtocol.readFieldZero();
            Map<String, String> headInfo = serverProtocol.getHead();

           String traceId = headInfo.get(TRACE_ID.getValue());
            String parentSpanId = headInfo.get(PARENT_SPAN_ID.getValue());
            String isSampled = headInfo.get(IS_SAMPLED.getValue());
            Boolean sampled = isSampled == null || Boolean.parseBoolean(isSampled);

            if (traceId != null && parentSpanId != null) {
                TraceUtils.startLocalTracer("rpc.thrift receive", traceId, parentSpanId, sampled);
                String methodName = tMessage.name;
                TraceUtils.submitAdditionalAnnotation(Constants.TRACE_THRIFT_METHOD, methodName);
                TTransport transport = in.getTransport();
                String hostAddress = ((TSocket) transport).getSocket().getRemoteSocketAddress().toString();
                TraceUtils.submitAdditionalAnnotation(Constants.TRACE_THRIFT_SERVER, hostAddress);
            }
            serverProtocol.resetTFramedTransport(in);

        }
        boolean result = realProcessor.process(in, out);
        if (in instanceof TTraceServerProtocol) {
            TraceUtils.endAndSendLocalTracer();
        }
        return result;
    }
}

使用方法

通过观察现有业务代码目前使用 thrift 的方式，总体来讲改动是有，但不大，用法可参看如下代码。

客户端

...
transport = new TSocket(SERVER_IP,SERVER_PORT,TIMEOUT);
TProtocol protocol = new TTraceClientProtocol(transport);
ThriftTestServcie.Client client = new ThriftTestServcie.Client(protocol);
transport.open();
...

服务端

... TServerSocket serverTransport = new TServerSocket(8081); TThreadPoolServer.Args ttpsArgs = new TThreadPoolServer.Args(serverTransport); ttpsArgs.processor(tTraceProcessor); ttpsArgs.protocolFactory(new TTraceServerProtocol.Factory()); TServer server = new TThreadPoolServer(ttpsArgs); server.serve(); ...