set client_timeout in header (#24)

XiaoMi · Dec 12, 2018 · 7bad957 · 7bad957
1 parent 90c7dc2
commit 7bad957
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 4 deletions.
diff --git a/src/main/java/com/xiaomi/infra/pegasus/operator/client_operator.java b/src/main/java/com/xiaomi/infra/pegasus/operator/client_operator.java
@@ -19,9 +19,10 @@ public client_operator(gpid gpid, String tableName) {
         this.rpc_error = new error_code();
     }
 
-    public final byte[] prepare_thrift_header(int body_length) {
+    public final byte[] prepare_thrift_header(int body_length, int client_timeout) {
         header.body_length = body_length;
         header.header_length = ThriftHeader.HEADER_LENGTH;
+        header.client_timeout = client_timeout;
         header.thread_hash = Tools.dsn_gpid_to_thread_hash(header.app_id, header.partition_index);
         header.partition_hash = 0;
         return header.toByteArray();

diff --git a/src/main/java/com/xiaomi/infra/pegasus/rpc/async/ReplicaSession.java b/src/main/java/com/xiaomi/infra/pegasus/rpc/async/ReplicaSession.java
@@ -34,6 +34,7 @@ public static class RequestEntry {
         public com.xiaomi.infra.pegasus.operator.client_operator op;
         public Runnable callback;
         public ScheduledFuture timeoutTask;
+        public long timeoutMs;
     }
 
     public enum ConnState {
@@ -83,6 +84,7 @@ public int asyncSend(client_operator op, Runnable callbackFunc, long timeoutInMi
         //the timer task is scheduled.
         pendingResponse.put(new Integer(entry.sequenceId), entry);
         entry.timeoutTask = addTimer(entry.sequenceId, timeoutInMilliseconds);
+        entry.timeoutMs = timeoutInMilliseconds;
 
         // We store the connection_state & netty channel in a struct so that they can fetch and update in atomic.
         // Moreover, we can avoid the lock protection when we want to get the netty channel for send message

diff --git a/src/main/java/com/xiaomi/infra/pegasus/rpc/async/ThriftFrameEncoder.java b/src/main/java/com/xiaomi/infra/pegasus/rpc/async/ThriftFrameEncoder.java
@@ -21,18 +21,19 @@ public ThriftFrameEncoder() {
 
     @Override
     protected ByteBuf allocateBuffer(ChannelHandlerContext ctx, ReplicaSession.RequestEntry entry, boolean preferDirect) throws Exception {
-        return preferDirect?ctx.alloc().ioBuffer(256):ctx.alloc().heapBuffer(256);
+        return preferDirect ? ctx.alloc().ioBuffer(256) : ctx.alloc().heapBuffer(256);
     }
 
     @Override
-    protected void encode(ChannelHandlerContext ctx, ReplicaSession.RequestEntry e, ByteBuf out) throws Exception{
+    protected void encode(ChannelHandlerContext ctx, ReplicaSession.RequestEntry e, ByteBuf out) throws Exception {
         int initIndex = out.writerIndex();
 
         // write the Memory buffer
         out.writerIndex(initIndex + ThriftHeader.HEADER_LENGTH);
         TBinaryProtocol protocol = new TBinaryProtocol(new TByteBufTransport(out));
         e.op.send_data(protocol, e.sequenceId);
-        out.setBytes(initIndex, e.op.prepare_thrift_header(out.readableBytes() - ThriftHeader.HEADER_LENGTH));
+        out.setBytes(initIndex, e.op.prepare_thrift_header(
+                out.readableBytes() - ThriftHeader.HEADER_LENGTH, (int) e.timeoutMs));
     }
 
     @Override