diff --git a/lib/go/test/tests/client_error_test.go b/lib/go/test/tests/client_error_test.go
index fdec4ea..a064163 100644
--- a/lib/go/test/tests/client_error_test.go
+++ b/lib/go/test/tests/client_error_test.go
@@ -37,84 +37,84 @@
 	if failAt == 0 {
 		err = failWith
 	}
-	last := protocol.EXPECT().WriteMessageBegin("testStruct", thrift.CALL, int32(1)).Return(err)
+	last := protocol.EXPECT().WriteMessageBegin(context.Background(), "testStruct", thrift.CALL, int32(1)).Return(err)
 	if failAt == 0 {
 		return true
 	}
 	if failAt == 1 {
 		err = failWith
 	}
-	last = protocol.EXPECT().WriteStructBegin("testStruct_args").Return(err).After(last)
+	last = protocol.EXPECT().WriteStructBegin(context.Background(), "testStruct_args").Return(err).After(last)
 	if failAt == 1 {
 		return true
 	}
 	if failAt == 2 {
 		err = failWith
 	}
-	last = protocol.EXPECT().WriteFieldBegin("thing", thrift.TType(thrift.STRUCT), int16(1)).Return(err).After(last)
+	last = protocol.EXPECT().WriteFieldBegin(context.Background(), "thing", thrift.TType(thrift.STRUCT), int16(1)).Return(err).After(last)
 	if failAt == 2 {
 		return true
 	}
 	if failAt == 3 {
 		err = failWith
 	}
-	last = protocol.EXPECT().WriteStructBegin("TestStruct").Return(err).After(last)
+	last = protocol.EXPECT().WriteStructBegin(context.Background(), "TestStruct").Return(err).After(last)
 	if failAt == 3 {
 		return true
 	}
 	if failAt == 4 {
 		err = failWith
 	}
-	last = protocol.EXPECT().WriteFieldBegin("m", thrift.TType(thrift.MAP), int16(1)).Return(err).After(last)
+	last = protocol.EXPECT().WriteFieldBegin(context.Background(), "m", thrift.TType(thrift.MAP), int16(1)).Return(err).After(last)
 	if failAt == 4 {
 		return true
 	}
 	if failAt == 5 {
 		err = failWith
 	}
-	last = protocol.EXPECT().WriteMapBegin(thrift.TType(thrift.STRING), thrift.TType(thrift.STRING), 0).Return(err).After(last)
+	last = protocol.EXPECT().WriteMapBegin(context.Background(), thrift.TType(thrift.STRING), thrift.TType(thrift.STRING), 0).Return(err).After(last)
 	if failAt == 5 {
 		return true
 	}
 	if failAt == 6 {
 		err = failWith
 	}
-	last = protocol.EXPECT().WriteMapEnd().Return(err).After(last)
+	last = protocol.EXPECT().WriteMapEnd(context.Background()).Return(err).After(last)
 	if failAt == 6 {
 		return true
 	}
 	if failAt == 7 {
 		err = failWith
 	}
-	last = protocol.EXPECT().WriteFieldEnd().Return(err).After(last)
+	last = protocol.EXPECT().WriteFieldEnd(context.Background()).Return(err).After(last)
 	if failAt == 7 {
 		return true
 	}
 	if failAt == 8 {
 		err = failWith
 	}
-	last = protocol.EXPECT().WriteFieldBegin("l", thrift.TType(thrift.LIST), int16(2)).Return(err).After(last)
+	last = protocol.EXPECT().WriteFieldBegin(context.Background(), "l", thrift.TType(thrift.LIST), int16(2)).Return(err).After(last)
 	if failAt == 8 {
 		return true
 	}
 	if failAt == 9 {
 		err = failWith
 	}
-	last = protocol.EXPECT().WriteListBegin(thrift.TType(thrift.STRING), 0).Return(err).After(last)
+	last = protocol.EXPECT().WriteListBegin(context.Background(), thrift.TType(thrift.STRING), 0).Return(err).After(last)
 	if failAt == 9 {
 		return true
 	}
 	if failAt == 10 {
 		err = failWith
 	}
-	last = protocol.EXPECT().WriteListEnd().Return(err).After(last)
+	last = protocol.EXPECT().WriteListEnd(context.Background()).Return(err).After(last)
 	if failAt == 10 {
 		return true
 	}
 	if failAt == 11 {
 		err = failWith
 	}
-	last = protocol.EXPECT().WriteFieldEnd().Return(err).After(last)
+	last = protocol.EXPECT().WriteFieldEnd(context.Background()).Return(err).After(last)
 	if failAt == 11 {
 		return true
 	}
@@ -122,91 +122,91 @@
 		err = failWith
 	}
 
-	last = protocol.EXPECT().WriteFieldBegin("s", thrift.TType(thrift.SET), int16(3)).Return(err).After(last)
+	last = protocol.EXPECT().WriteFieldBegin(context.Background(), "s", thrift.TType(thrift.SET), int16(3)).Return(err).After(last)
 	if failAt == 12 {
 		return true
 	}
 	if failAt == 13 {
 		err = failWith
 	}
-	last = protocol.EXPECT().WriteSetBegin(thrift.TType(thrift.STRING), 0).Return(err).After(last)
+	last = protocol.EXPECT().WriteSetBegin(context.Background(), thrift.TType(thrift.STRING), 0).Return(err).After(last)
 	if failAt == 13 {
 		return true
 	}
 	if failAt == 14 {
 		err = failWith
 	}
-	last = protocol.EXPECT().WriteSetEnd().Return(err).After(last)
+	last = protocol.EXPECT().WriteSetEnd(context.Background()).Return(err).After(last)
 	if failAt == 14 {
 		return true
 	}
 	if failAt == 15 {
 		err = failWith
 	}
-	last = protocol.EXPECT().WriteFieldEnd().Return(err).After(last)
+	last = protocol.EXPECT().WriteFieldEnd(context.Background()).Return(err).After(last)
 	if failAt == 15 {
 		return true
 	}
 	if failAt == 16 {
 		err = failWith
 	}
-	last = protocol.EXPECT().WriteFieldBegin("i", thrift.TType(thrift.I32), int16(4)).Return(err).After(last)
+	last = protocol.EXPECT().WriteFieldBegin(context.Background(), "i", thrift.TType(thrift.I32), int16(4)).Return(err).After(last)
 	if failAt == 16 {
 		return true
 	}
 	if failAt == 17 {
 		err = failWith
 	}
-	last = protocol.EXPECT().WriteI32(int32(3)).Return(err).After(last)
+	last = protocol.EXPECT().WriteI32(context.Background(), int32(3)).Return(err).After(last)
 	if failAt == 17 {
 		return true
 	}
 	if failAt == 18 {
 		err = failWith
 	}
-	last = protocol.EXPECT().WriteFieldEnd().Return(err).After(last)
+	last = protocol.EXPECT().WriteFieldEnd(context.Background()).Return(err).After(last)
 	if failAt == 18 {
 		return true
 	}
 	if failAt == 19 {
 		err = failWith
 	}
-	last = protocol.EXPECT().WriteFieldStop().Return(err).After(last)
+	last = protocol.EXPECT().WriteFieldStop(context.Background()).Return(err).After(last)
 	if failAt == 19 {
 		return true
 	}
 	if failAt == 20 {
 		err = failWith
 	}
-	last = protocol.EXPECT().WriteStructEnd().Return(err).After(last)
+	last = protocol.EXPECT().WriteStructEnd(context.Background()).Return(err).After(last)
 	if failAt == 20 {
 		return true
 	}
 	if failAt == 21 {
 		err = failWith
 	}
-	last = protocol.EXPECT().WriteFieldEnd().Return(err).After(last)
+	last = protocol.EXPECT().WriteFieldEnd(context.Background()).Return(err).After(last)
 	if failAt == 21 {
 		return true
 	}
 	if failAt == 22 {
 		err = failWith
 	}
-	last = protocol.EXPECT().WriteFieldStop().Return(err).After(last)
+	last = protocol.EXPECT().WriteFieldStop(context.Background()).Return(err).After(last)
 	if failAt == 22 {
 		return true
 	}
 	if failAt == 23 {
 		err = failWith
 	}
-	last = protocol.EXPECT().WriteStructEnd().Return(err).After(last)
+	last = protocol.EXPECT().WriteStructEnd(context.Background()).Return(err).After(last)
 	if failAt == 23 {
 		return true
 	}
 	if failAt == 24 {
 		err = failWith
 	}
-	last = protocol.EXPECT().WriteMessageEnd().Return(err).After(last)
+	last = protocol.EXPECT().WriteMessageEnd(context.Background()).Return(err).After(last)
 	if failAt == 24 {
 		return true
 	}
@@ -220,175 +220,175 @@
 	if failAt == 26 {
 		err = failWith
 	}
-	last = protocol.EXPECT().ReadMessageBegin().Return("testStruct", thrift.REPLY, int32(1), err).After(last)
+	last = protocol.EXPECT().ReadMessageBegin(context.Background()).Return("testStruct", thrift.REPLY, int32(1), err).After(last)
 	if failAt == 26 {
 		return true
 	}
 	if failAt == 27 {
 		err = failWith
 	}
-	last = protocol.EXPECT().ReadStructBegin().Return("testStruct_args", err).After(last)
+	last = protocol.EXPECT().ReadStructBegin(context.Background()).Return("testStruct_args", err).After(last)
 	if failAt == 27 {
 		return true
 	}
 	if failAt == 28 {
 		err = failWith
 	}
-	last = protocol.EXPECT().ReadFieldBegin().Return("_", thrift.TType(thrift.STRUCT), int16(0), err).After(last)
+	last = protocol.EXPECT().ReadFieldBegin(context.Background()).Return("_", thrift.TType(thrift.STRUCT), int16(0), err).After(last)
 	if failAt == 28 {
 		return true
 	}
 	if failAt == 29 {
 		err = failWith
 	}
-	last = protocol.EXPECT().ReadStructBegin().Return("TestStruct", err).After(last)
+	last = protocol.EXPECT().ReadStructBegin(context.Background()).Return("TestStruct", err).After(last)
 	if failAt == 29 {
 		return true
 	}
 	if failAt == 30 {
 		err = failWith
 	}
-	last = protocol.EXPECT().ReadFieldBegin().Return("m", thrift.TType(thrift.MAP), int16(1), err).After(last)
+	last = protocol.EXPECT().ReadFieldBegin(context.Background()).Return("m", thrift.TType(thrift.MAP), int16(1), err).After(last)
 	if failAt == 30 {
 		return true
 	}
 	if failAt == 31 {
 		err = failWith
 	}
-	last = protocol.EXPECT().ReadMapBegin().Return(thrift.TType(thrift.STRING), thrift.TType(thrift.STRING), 0, err).After(last)
+	last = protocol.EXPECT().ReadMapBegin(context.Background()).Return(thrift.TType(thrift.STRING), thrift.TType(thrift.STRING), 0, err).After(last)
 	if failAt == 31 {
 		return true
 	}
 	if failAt == 32 {
 		err = failWith
 	}
-	last = protocol.EXPECT().ReadMapEnd().Return(err).After(last)
+	last = protocol.EXPECT().ReadMapEnd(context.Background()).Return(err).After(last)
 	if failAt == 32 {
 		return true
 	}
 	if failAt == 33 {
 		err = failWith
 	}
-	last = protocol.EXPECT().ReadFieldEnd().Return(err).After(last)
+	last = protocol.EXPECT().ReadFieldEnd(context.Background()).Return(err).After(last)
 	if failAt == 33 {
 		return true
 	}
 	if failAt == 34 {
 		err = failWith
 	}
-	last = protocol.EXPECT().ReadFieldBegin().Return("l", thrift.TType(thrift.LIST), int16(2), err).After(last)
+	last = protocol.EXPECT().ReadFieldBegin(context.Background()).Return("l", thrift.TType(thrift.LIST), int16(2), err).After(last)
 	if failAt == 34 {
 		return true
 	}
 	if failAt == 35 {
 		err = failWith
 	}
-	last = protocol.EXPECT().ReadListBegin().Return(thrift.TType(thrift.STRING), 0, err).After(last)
+	last = protocol.EXPECT().ReadListBegin(context.Background()).Return(thrift.TType(thrift.STRING), 0, err).After(last)
 	if failAt == 35 {
 		return true
 	}
 	if failAt == 36 {
 		err = failWith
 	}
-	last = protocol.EXPECT().ReadListEnd().Return(err).After(last)
+	last = protocol.EXPECT().ReadListEnd(context.Background()).Return(err).After(last)
 	if failAt == 36 {
 		return true
 	}
 	if failAt == 37 {
 		err = failWith
 	}
-	last = protocol.EXPECT().ReadFieldEnd().Return(err).After(last)
+	last = protocol.EXPECT().ReadFieldEnd(context.Background()).Return(err).After(last)
 	if failAt == 37 {
 		return true
 	}
 	if failAt == 38 {
 		err = failWith
 	}
-	last = protocol.EXPECT().ReadFieldBegin().Return("s", thrift.TType(thrift.SET), int16(3), err).After(last)
+	last = protocol.EXPECT().ReadFieldBegin(context.Background()).Return("s", thrift.TType(thrift.SET), int16(3), err).After(last)
 	if failAt == 38 {
 		return true
 	}
 	if failAt == 39 {
 		err = failWith
 	}
-	last = protocol.EXPECT().ReadSetBegin().Return(thrift.TType(thrift.STRING), 0, err).After(last)
+	last = protocol.EXPECT().ReadSetBegin(context.Background()).Return(thrift.TType(thrift.STRING), 0, err).After(last)
 	if failAt == 39 {
 		return true
 	}
 	if failAt == 40 {
 		err = failWith
 	}
-	last = protocol.EXPECT().ReadSetEnd().Return(err).After(last)
+	last = protocol.EXPECT().ReadSetEnd(context.Background()).Return(err).After(last)
 	if failAt == 40 {
 		return true
 	}
 	if failAt == 41 {
 		err = failWith
 	}
-	last = protocol.EXPECT().ReadFieldEnd().Return(err).After(last)
+	last = protocol.EXPECT().ReadFieldEnd(context.Background()).Return(err).After(last)
 	if failAt == 41 {
 		return true
 	}
 	if failAt == 42 {
 		err = failWith
 	}
-	last = protocol.EXPECT().ReadFieldBegin().Return("i", thrift.TType(thrift.I32), int16(4), err).After(last)
+	last = protocol.EXPECT().ReadFieldBegin(context.Background()).Return("i", thrift.TType(thrift.I32), int16(4), err).After(last)
 	if failAt == 42 {
 		return true
 	}
 	if failAt == 43 {
 		err = failWith
 	}
-	last = protocol.EXPECT().ReadI32().Return(int32(3), err).After(last)
+	last = protocol.EXPECT().ReadI32(context.Background()).Return(int32(3), err).After(last)
 	if failAt == 43 {
 		return true
 	}
 	if failAt == 44 {
 		err = failWith
 	}
-	last = protocol.EXPECT().ReadFieldEnd().Return(err).After(last)
+	last = protocol.EXPECT().ReadFieldEnd(context.Background()).Return(err).After(last)
 	if failAt == 44 {
 		return true
 	}
 	if failAt == 45 {
 		err = failWith
 	}
-	last = protocol.EXPECT().ReadFieldBegin().Return("_", thrift.TType(thrift.STOP), int16(5), err).After(last)
+	last = protocol.EXPECT().ReadFieldBegin(context.Background()).Return("_", thrift.TType(thrift.STOP), int16(5), err).After(last)
 	if failAt == 45 {
 		return true
 	}
 	if failAt == 46 {
 		err = failWith
 	}
-	last = protocol.EXPECT().ReadStructEnd().Return(err).After(last)
+	last = protocol.EXPECT().ReadStructEnd(context.Background()).Return(err).After(last)
 	if failAt == 46 {
 		return true
 	}
 	if failAt == 47 {
 		err = failWith
 	}
-	last = protocol.EXPECT().ReadFieldEnd().Return(err).After(last)
+	last = protocol.EXPECT().ReadFieldEnd(context.Background()).Return(err).After(last)
 	if failAt == 47 {
 		return true
 	}
 	if failAt == 48 {
 		err = failWith
 	}
-	last = protocol.EXPECT().ReadFieldBegin().Return("_", thrift.TType(thrift.STOP), int16(1), err).After(last)
+	last = protocol.EXPECT().ReadFieldBegin(context.Background()).Return("_", thrift.TType(thrift.STOP), int16(1), err).After(last)
 	if failAt == 48 {
 		return true
 	}
 	if failAt == 49 {
 		err = failWith
 	}
-	last = protocol.EXPECT().ReadStructEnd().Return(err).After(last)
+	last = protocol.EXPECT().ReadStructEnd(context.Background()).Return(err).After(last)
 	if failAt == 49 {
 		return true
 	}
 	if failAt == 50 {
 		err = failWith
 	}
-	last = protocol.EXPECT().ReadMessageEnd().Return(err).After(last)
+	last = protocol.EXPECT().ReadMessageEnd(context.Background()).Return(err).After(last)
 	if failAt == 50 {
 		return true
 	}
@@ -529,91 +529,91 @@
 	var err error = nil
 
 	// No need to test failure in this block, because it is covered in other test cases
-	last := protocol.EXPECT().WriteMessageBegin("testString", thrift.CALL, int32(1))
-	last = protocol.EXPECT().WriteStructBegin("testString_args").After(last)
-	last = protocol.EXPECT().WriteFieldBegin("s", thrift.TType(thrift.STRING), int16(1)).After(last)
-	last = protocol.EXPECT().WriteString("test").After(last)
-	last = protocol.EXPECT().WriteFieldEnd().After(last)
-	last = protocol.EXPECT().WriteFieldStop().After(last)
-	last = protocol.EXPECT().WriteStructEnd().After(last)
-	last = protocol.EXPECT().WriteMessageEnd().After(last)
+	last := protocol.EXPECT().WriteMessageBegin(context.Background(), "testString", thrift.CALL, int32(1))
+	last = protocol.EXPECT().WriteStructBegin(context.Background(), "testString_args").After(last)
+	last = protocol.EXPECT().WriteFieldBegin(context.Background(), "s", thrift.TType(thrift.STRING), int16(1)).After(last)
+	last = protocol.EXPECT().WriteString(context.Background(), "test").After(last)
+	last = protocol.EXPECT().WriteFieldEnd(context.Background()).After(last)
+	last = protocol.EXPECT().WriteFieldStop(context.Background()).After(last)
+	last = protocol.EXPECT().WriteStructEnd(context.Background()).After(last)
+	last = protocol.EXPECT().WriteMessageEnd(context.Background()).After(last)
 	last = protocol.EXPECT().Flush(context.Background()).After(last)
 
 	// Reading the exception, might fail.
 	if failAt == 0 {
 		err = failWith
 	}
-	last = protocol.EXPECT().ReadMessageBegin().Return("testString", thrift.EXCEPTION, int32(1), err).After(last)
+	last = protocol.EXPECT().ReadMessageBegin(context.Background()).Return("testString", thrift.EXCEPTION, int32(1), err).After(last)
 	if failAt == 0 {
 		return true
 	}
 	if failAt == 1 {
 		err = failWith
 	}
-	last = protocol.EXPECT().ReadStructBegin().Return("TApplicationException", err).After(last)
+	last = protocol.EXPECT().ReadStructBegin(context.Background()).Return("TApplicationException", err).After(last)
 	if failAt == 1 {
 		return true
 	}
 	if failAt == 2 {
 		err = failWith
 	}
-	last = protocol.EXPECT().ReadFieldBegin().Return("message", thrift.TType(thrift.STRING), int16(1), err).After(last)
+	last = protocol.EXPECT().ReadFieldBegin(context.Background()).Return("message", thrift.TType(thrift.STRING), int16(1), err).After(last)
 	if failAt == 2 {
 		return true
 	}
 	if failAt == 3 {
 		err = failWith
 	}
-	last = protocol.EXPECT().ReadString().Return("test", err).After(last)
+	last = protocol.EXPECT().ReadString(context.Background()).Return("test", err).After(last)
 	if failAt == 3 {
 		return true
 	}
 	if failAt == 4 {
 		err = failWith
 	}
-	last = protocol.EXPECT().ReadFieldEnd().Return(err).After(last)
+	last = protocol.EXPECT().ReadFieldEnd(context.Background()).Return(err).After(last)
 	if failAt == 4 {
 		return true
 	}
 	if failAt == 5 {
 		err = failWith
 	}
-	last = protocol.EXPECT().ReadFieldBegin().Return("type", thrift.TType(thrift.I32), int16(2), err).After(last)
+	last = protocol.EXPECT().ReadFieldBegin(context.Background()).Return("type", thrift.TType(thrift.I32), int16(2), err).After(last)
 	if failAt == 5 {
 		return true
 	}
 	if failAt == 6 {
 		err = failWith
 	}
-	last = protocol.EXPECT().ReadI32().Return(int32(thrift.PROTOCOL_ERROR), err).After(last)
+	last = protocol.EXPECT().ReadI32(context.Background()).Return(int32(thrift.PROTOCOL_ERROR), err).After(last)
 	if failAt == 6 {
 		return true
 	}
 	if failAt == 7 {
 		err = failWith
 	}
-	last = protocol.EXPECT().ReadFieldEnd().Return(err).After(last)
+	last = protocol.EXPECT().ReadFieldEnd(context.Background()).Return(err).After(last)
 	if failAt == 7 {
 		return true
 	}
 	if failAt == 8 {
 		err = failWith
 	}
-	last = protocol.EXPECT().ReadFieldBegin().Return("_", thrift.TType(thrift.STOP), int16(2), err).After(last)
+	last = protocol.EXPECT().ReadFieldBegin(context.Background()).Return("_", thrift.TType(thrift.STOP), int16(2), err).After(last)
 	if failAt == 8 {
 		return true
 	}
 	if failAt == 9 {
 		err = failWith
 	}
-	last = protocol.EXPECT().ReadStructEnd().Return(err).After(last)
+	last = protocol.EXPECT().ReadStructEnd(context.Background()).Return(err).After(last)
 	if failAt == 9 {
 		return true
 	}
 	if failAt == 10 {
 		err = failWith
 	}
-	last = protocol.EXPECT().ReadMessageEnd().Return(err).After(last)
+	last = protocol.EXPECT().ReadMessageEnd(context.Background()).Return(err).After(last)
 	if failAt == 10 {
 		return true
 	}
@@ -697,16 +697,16 @@
 	mockCtrl := gomock.NewController(t)
 	protocol := NewMockTProtocol(mockCtrl)
 	gomock.InOrder(
-		protocol.EXPECT().WriteMessageBegin("testString", thrift.CALL, int32(1)),
-		protocol.EXPECT().WriteStructBegin("testString_args"),
-		protocol.EXPECT().WriteFieldBegin("s", thrift.TType(thrift.STRING), int16(1)),
-		protocol.EXPECT().WriteString("test"),
-		protocol.EXPECT().WriteFieldEnd(),
-		protocol.EXPECT().WriteFieldStop(),
-		protocol.EXPECT().WriteStructEnd(),
-		protocol.EXPECT().WriteMessageEnd(),
+		protocol.EXPECT().WriteMessageBegin(context.Background(), "testString", thrift.CALL, int32(1)),
+		protocol.EXPECT().WriteStructBegin(context.Background(), "testString_args"),
+		protocol.EXPECT().WriteFieldBegin(context.Background(), "s", thrift.TType(thrift.STRING), int16(1)),
+		protocol.EXPECT().WriteString(context.Background(), "test"),
+		protocol.EXPECT().WriteFieldEnd(context.Background()),
+		protocol.EXPECT().WriteFieldStop(context.Background()),
+		protocol.EXPECT().WriteStructEnd(context.Background()),
+		protocol.EXPECT().WriteMessageEnd(context.Background()),
 		protocol.EXPECT().Flush(context.Background()),
-		protocol.EXPECT().ReadMessageBegin().Return("testString", thrift.REPLY, int32(2), nil),
+		protocol.EXPECT().ReadMessageBegin(context.Background()).Return("testString", thrift.REPLY, int32(2), nil),
 	)
 
 	client := errortest.NewErrorTestClient(thrift.NewTStandardClient(protocol, protocol))
@@ -728,16 +728,16 @@
 	transport := thrift.NewTMemoryBuffer()
 	protocol := NewMockTProtocol(mockCtrl)
 	gomock.InOrder(
-		protocol.EXPECT().WriteMessageBegin("testString", thrift.CALL, int32(1)),
-		protocol.EXPECT().WriteStructBegin("testString_args"),
-		protocol.EXPECT().WriteFieldBegin("s", thrift.TType(thrift.STRING), int16(1)),
-		protocol.EXPECT().WriteString("test"),
-		protocol.EXPECT().WriteFieldEnd(),
-		protocol.EXPECT().WriteFieldStop(),
-		protocol.EXPECT().WriteStructEnd(),
-		protocol.EXPECT().WriteMessageEnd(),
+		protocol.EXPECT().WriteMessageBegin(context.Background(), "testString", thrift.CALL, int32(1)),
+		protocol.EXPECT().WriteStructBegin(context.Background(), "testString_args"),
+		protocol.EXPECT().WriteFieldBegin(context.Background(), "s", thrift.TType(thrift.STRING), int16(1)),
+		protocol.EXPECT().WriteString(context.Background(), "test"),
+		protocol.EXPECT().WriteFieldEnd(context.Background()),
+		protocol.EXPECT().WriteFieldStop(context.Background()),
+		protocol.EXPECT().WriteStructEnd(context.Background()),
+		protocol.EXPECT().WriteMessageEnd(context.Background()),
 		protocol.EXPECT().Flush(context.Background()),
-		protocol.EXPECT().ReadMessageBegin().Return("testString", thrift.REPLY, int32(2), nil),
+		protocol.EXPECT().ReadMessageBegin(context.Background()).Return("testString", thrift.REPLY, int32(2), nil),
 	)
 
 	client := errortest.NewErrorTestClientProtocol(transport, protocol, protocol)
@@ -757,16 +757,16 @@
 	mockCtrl := gomock.NewController(t)
 	protocol := NewMockTProtocol(mockCtrl)
 	gomock.InOrder(
-		protocol.EXPECT().WriteMessageBegin("testString", thrift.CALL, int32(1)),
-		protocol.EXPECT().WriteStructBegin("testString_args"),
-		protocol.EXPECT().WriteFieldBegin("s", thrift.TType(thrift.STRING), int16(1)),
-		protocol.EXPECT().WriteString("test"),
-		protocol.EXPECT().WriteFieldEnd(),
-		protocol.EXPECT().WriteFieldStop(),
-		protocol.EXPECT().WriteStructEnd(),
-		protocol.EXPECT().WriteMessageEnd(),
+		protocol.EXPECT().WriteMessageBegin(context.Background(), "testString", thrift.CALL, int32(1)),
+		protocol.EXPECT().WriteStructBegin(context.Background(), "testString_args"),
+		protocol.EXPECT().WriteFieldBegin(context.Background(), "s", thrift.TType(thrift.STRING), int16(1)),
+		protocol.EXPECT().WriteString(context.Background(), "test"),
+		protocol.EXPECT().WriteFieldEnd(context.Background()),
+		protocol.EXPECT().WriteFieldStop(context.Background()),
+		protocol.EXPECT().WriteStructEnd(context.Background()),
+		protocol.EXPECT().WriteMessageEnd(context.Background()),
 		protocol.EXPECT().Flush(context.Background()),
-		protocol.EXPECT().ReadMessageBegin().Return("unknown", thrift.REPLY, int32(1), nil),
+		protocol.EXPECT().ReadMessageBegin(context.Background()).Return("unknown", thrift.REPLY, int32(1), nil),
 	)
 
 	client := errortest.NewErrorTestClient(thrift.NewTStandardClient(protocol, protocol))
@@ -788,16 +788,16 @@
 	transport := thrift.NewTMemoryBuffer()
 	protocol := NewMockTProtocol(mockCtrl)
 	gomock.InOrder(
-		protocol.EXPECT().WriteMessageBegin("testString", thrift.CALL, int32(1)),
-		protocol.EXPECT().WriteStructBegin("testString_args"),
-		protocol.EXPECT().WriteFieldBegin("s", thrift.TType(thrift.STRING), int16(1)),
-		protocol.EXPECT().WriteString("test"),
-		protocol.EXPECT().WriteFieldEnd(),
-		protocol.EXPECT().WriteFieldStop(),
-		protocol.EXPECT().WriteStructEnd(),
-		protocol.EXPECT().WriteMessageEnd(),
+		protocol.EXPECT().WriteMessageBegin(context.Background(), "testString", thrift.CALL, int32(1)),
+		protocol.EXPECT().WriteStructBegin(context.Background(), "testString_args"),
+		protocol.EXPECT().WriteFieldBegin(context.Background(), "s", thrift.TType(thrift.STRING), int16(1)),
+		protocol.EXPECT().WriteString(context.Background(), "test"),
+		protocol.EXPECT().WriteFieldEnd(context.Background()),
+		protocol.EXPECT().WriteFieldStop(context.Background()),
+		protocol.EXPECT().WriteStructEnd(context.Background()),
+		protocol.EXPECT().WriteMessageEnd(context.Background()),
 		protocol.EXPECT().Flush(context.Background()),
-		protocol.EXPECT().ReadMessageBegin().Return("unknown", thrift.REPLY, int32(1), nil),
+		protocol.EXPECT().ReadMessageBegin(context.Background()).Return("unknown", thrift.REPLY, int32(1), nil),
 	)
 
 	client := errortest.NewErrorTestClientProtocol(transport, protocol, protocol)
@@ -817,16 +817,16 @@
 	mockCtrl := gomock.NewController(t)
 	protocol := NewMockTProtocol(mockCtrl)
 	gomock.InOrder(
-		protocol.EXPECT().WriteMessageBegin("testString", thrift.CALL, int32(1)),
-		protocol.EXPECT().WriteStructBegin("testString_args"),
-		protocol.EXPECT().WriteFieldBegin("s", thrift.TType(thrift.STRING), int16(1)),
-		protocol.EXPECT().WriteString("test"),
-		protocol.EXPECT().WriteFieldEnd(),
-		protocol.EXPECT().WriteFieldStop(),
-		protocol.EXPECT().WriteStructEnd(),
-		protocol.EXPECT().WriteMessageEnd(),
+		protocol.EXPECT().WriteMessageBegin(context.Background(), "testString", thrift.CALL, int32(1)),
+		protocol.EXPECT().WriteStructBegin(context.Background(), "testString_args"),
+		protocol.EXPECT().WriteFieldBegin(context.Background(), "s", thrift.TType(thrift.STRING), int16(1)),
+		protocol.EXPECT().WriteString(context.Background(), "test"),
+		protocol.EXPECT().WriteFieldEnd(context.Background()),
+		protocol.EXPECT().WriteFieldStop(context.Background()),
+		protocol.EXPECT().WriteStructEnd(context.Background()),
+		protocol.EXPECT().WriteMessageEnd(context.Background()),
 		protocol.EXPECT().Flush(context.Background()),
-		protocol.EXPECT().ReadMessageBegin().Return("testString", thrift.INVALID_TMESSAGE_TYPE, int32(1), nil),
+		protocol.EXPECT().ReadMessageBegin(context.Background()).Return("testString", thrift.INVALID_TMESSAGE_TYPE, int32(1), nil),
 	)
 
 	client := errortest.NewErrorTestClient(thrift.NewTStandardClient(protocol, protocol))
@@ -848,16 +848,16 @@
 	transport := thrift.NewTMemoryBuffer()
 	protocol := NewMockTProtocol(mockCtrl)
 	gomock.InOrder(
-		protocol.EXPECT().WriteMessageBegin("testString", thrift.CALL, int32(1)),
-		protocol.EXPECT().WriteStructBegin("testString_args"),
-		protocol.EXPECT().WriteFieldBegin("s", thrift.TType(thrift.STRING), int16(1)),
-		protocol.EXPECT().WriteString("test"),
-		protocol.EXPECT().WriteFieldEnd(),
-		protocol.EXPECT().WriteFieldStop(),
-		protocol.EXPECT().WriteStructEnd(),
-		protocol.EXPECT().WriteMessageEnd(),
+		protocol.EXPECT().WriteMessageBegin(context.Background(), "testString", thrift.CALL, int32(1)),
+		protocol.EXPECT().WriteStructBegin(context.Background(), "testString_args"),
+		protocol.EXPECT().WriteFieldBegin(context.Background(), "s", thrift.TType(thrift.STRING), int16(1)),
+		protocol.EXPECT().WriteString(context.Background(), "test"),
+		protocol.EXPECT().WriteFieldEnd(context.Background()),
+		protocol.EXPECT().WriteFieldStop(context.Background()),
+		protocol.EXPECT().WriteStructEnd(context.Background()),
+		protocol.EXPECT().WriteMessageEnd(context.Background()),
 		protocol.EXPECT().Flush(context.Background()),
-		protocol.EXPECT().ReadMessageBegin().Return("testString", thrift.INVALID_TMESSAGE_TYPE, int32(1), nil),
+		protocol.EXPECT().ReadMessageBegin(context.Background()).Return("testString", thrift.INVALID_TMESSAGE_TYPE, int32(1), nil),
 	)
 
 	client := errortest.NewErrorTestClientProtocol(transport, protocol, protocol)
diff --git a/lib/go/test/tests/optional_fields_test.go b/lib/go/test/tests/optional_fields_test.go
index 34ad660..7e240e6 100644
--- a/lib/go/test/tests/optional_fields_test.go
+++ b/lib/go/test/tests/optional_fields_test.go
@@ -21,6 +21,7 @@
 
 import (
 	"bytes"
+	"context"
 	gomock "github.com/golang/mock/gomock"
 	"optionalfieldstest"
 	"testing"
@@ -185,12 +186,12 @@
 	defer mockCtrl.Finish()
 	proto := NewMockTProtocol(mockCtrl)
 	gomock.InOrder(
-		proto.EXPECT().WriteStructBegin("all_optional").Return(nil),
-		proto.EXPECT().WriteFieldStop().Return(nil),
-		proto.EXPECT().WriteStructEnd().Return(nil),
+		proto.EXPECT().WriteStructBegin(context.Background(), "all_optional").Return(nil),
+		proto.EXPECT().WriteFieldStop(context.Background()).Return(nil),
+		proto.EXPECT().WriteStructEnd(context.Background()).Return(nil),
 	)
 	ao := optionalfieldstest.NewAllOptional()
-	ao.Write(proto)
+	ao.Write(context.Background(), proto)
 }
 
 func TestNoSetToDefaultFieldsOnWire(t *testing.T) {
@@ -198,13 +199,13 @@
 	defer mockCtrl.Finish()
 	proto := NewMockTProtocol(mockCtrl)
 	gomock.InOrder(
-		proto.EXPECT().WriteStructBegin("all_optional").Return(nil),
-		proto.EXPECT().WriteFieldStop().Return(nil),
-		proto.EXPECT().WriteStructEnd().Return(nil),
+		proto.EXPECT().WriteStructBegin(context.Background(), "all_optional").Return(nil),
+		proto.EXPECT().WriteFieldStop(context.Background()).Return(nil),
+		proto.EXPECT().WriteStructEnd(context.Background()).Return(nil),
 	)
 	ao := optionalfieldstest.NewAllOptional()
 	ao.I = 42
-	ao.Write(proto)
+	ao.Write(context.Background(), proto)
 }
 
 //Make sure that only one field is being serialized when set to non-default
@@ -213,16 +214,16 @@
 	defer mockCtrl.Finish()
 	proto := NewMockTProtocol(mockCtrl)
 	gomock.InOrder(
-		proto.EXPECT().WriteStructBegin("all_optional").Return(nil),
-		proto.EXPECT().WriteFieldBegin("i", thrift.TType(thrift.I64), int16(2)).Return(nil),
-		proto.EXPECT().WriteI64(int64(123)).Return(nil),
-		proto.EXPECT().WriteFieldEnd().Return(nil),
-		proto.EXPECT().WriteFieldStop().Return(nil),
-		proto.EXPECT().WriteStructEnd().Return(nil),
+		proto.EXPECT().WriteStructBegin(context.Background(), "all_optional").Return(nil),
+		proto.EXPECT().WriteFieldBegin(context.Background(), "i", thrift.TType(thrift.I64), int16(2)).Return(nil),
+		proto.EXPECT().WriteI64(context.Background(), int64(123)).Return(nil),
+		proto.EXPECT().WriteFieldEnd(context.Background()).Return(nil),
+		proto.EXPECT().WriteFieldStop(context.Background()).Return(nil),
+		proto.EXPECT().WriteStructEnd(context.Background()).Return(nil),
 	)
 	ao := optionalfieldstest.NewAllOptional()
 	ao.I = 123
-	ao.Write(proto)
+	ao.Write(context.Background(), proto)
 }
 
 func TestOneLSetFieldOnWire(t *testing.T) {
@@ -230,19 +231,19 @@
 	defer mockCtrl.Finish()
 	proto := NewMockTProtocol(mockCtrl)
 	gomock.InOrder(
-		proto.EXPECT().WriteStructBegin("all_optional").Return(nil),
-		proto.EXPECT().WriteFieldBegin("l", thrift.TType(thrift.LIST), int16(9)).Return(nil),
-		proto.EXPECT().WriteListBegin(thrift.TType(thrift.I64), 2).Return(nil),
-		proto.EXPECT().WriteI64(int64(1)).Return(nil),
-		proto.EXPECT().WriteI64(int64(2)).Return(nil),
-		proto.EXPECT().WriteListEnd().Return(nil),
-		proto.EXPECT().WriteFieldEnd().Return(nil),
-		proto.EXPECT().WriteFieldStop().Return(nil),
-		proto.EXPECT().WriteStructEnd().Return(nil),
+		proto.EXPECT().WriteStructBegin(context.Background(), "all_optional").Return(nil),
+		proto.EXPECT().WriteFieldBegin(context.Background(), "l", thrift.TType(thrift.LIST), int16(9)).Return(nil),
+		proto.EXPECT().WriteListBegin(context.Background(), thrift.TType(thrift.I64), 2).Return(nil),
+		proto.EXPECT().WriteI64(context.Background(), int64(1)).Return(nil),
+		proto.EXPECT().WriteI64(context.Background(), int64(2)).Return(nil),
+		proto.EXPECT().WriteListEnd(context.Background()).Return(nil),
+		proto.EXPECT().WriteFieldEnd(context.Background()).Return(nil),
+		proto.EXPECT().WriteFieldStop(context.Background()).Return(nil),
+		proto.EXPECT().WriteStructEnd(context.Background()).Return(nil),
 	)
 	ao := optionalfieldstest.NewAllOptional()
 	ao.L = []int64{1, 2}
-	ao.Write(proto)
+	ao.Write(context.Background(), proto)
 }
 
 func TestOneBinSetFieldOnWire(t *testing.T) {
@@ -250,16 +251,16 @@
 	defer mockCtrl.Finish()
 	proto := NewMockTProtocol(mockCtrl)
 	gomock.InOrder(
-		proto.EXPECT().WriteStructBegin("all_optional").Return(nil),
-		proto.EXPECT().WriteFieldBegin("bin", thrift.TType(thrift.STRING), int16(13)).Return(nil),
-		proto.EXPECT().WriteBinary([]byte("somebytestring")).Return(nil),
-		proto.EXPECT().WriteFieldEnd().Return(nil),
-		proto.EXPECT().WriteFieldStop().Return(nil),
-		proto.EXPECT().WriteStructEnd().Return(nil),
+		proto.EXPECT().WriteStructBegin(context.Background(), "all_optional").Return(nil),
+		proto.EXPECT().WriteFieldBegin(context.Background(), "bin", thrift.TType(thrift.STRING), int16(13)).Return(nil),
+		proto.EXPECT().WriteBinary(context.Background(), []byte("somebytestring")).Return(nil),
+		proto.EXPECT().WriteFieldEnd(context.Background()).Return(nil),
+		proto.EXPECT().WriteFieldStop(context.Background()).Return(nil),
+		proto.EXPECT().WriteStructEnd(context.Background()).Return(nil),
 	)
 	ao := optionalfieldstest.NewAllOptional()
 	ao.Bin = []byte("somebytestring")
-	ao.Write(proto)
+	ao.Write(context.Background(), proto)
 }
 
 func TestOneEmptyBinSetFieldOnWire(t *testing.T) {
@@ -267,14 +268,14 @@
 	defer mockCtrl.Finish()
 	proto := NewMockTProtocol(mockCtrl)
 	gomock.InOrder(
-		proto.EXPECT().WriteStructBegin("all_optional").Return(nil),
-		proto.EXPECT().WriteFieldBegin("bin", thrift.TType(thrift.STRING), int16(13)).Return(nil),
-		proto.EXPECT().WriteBinary([]byte{}).Return(nil),
-		proto.EXPECT().WriteFieldEnd().Return(nil),
-		proto.EXPECT().WriteFieldStop().Return(nil),
-		proto.EXPECT().WriteStructEnd().Return(nil),
+		proto.EXPECT().WriteStructBegin(context.Background(), "all_optional").Return(nil),
+		proto.EXPECT().WriteFieldBegin(context.Background(), "bin", thrift.TType(thrift.STRING), int16(13)).Return(nil),
+		proto.EXPECT().WriteBinary(context.Background(), []byte{}).Return(nil),
+		proto.EXPECT().WriteFieldEnd(context.Background()).Return(nil),
+		proto.EXPECT().WriteFieldStop(context.Background()).Return(nil),
+		proto.EXPECT().WriteStructEnd(context.Background()).Return(nil),
 	)
 	ao := optionalfieldstest.NewAllOptional()
 	ao.Bin = []byte{}
-	ao.Write(proto)
+	ao.Write(context.Background(), proto)
 }
diff --git a/lib/go/test/tests/protocol_mock.go b/lib/go/test/tests/protocol_mock.go
index 51d7a02..793e4e1 100644
--- a/lib/go/test/tests/protocol_mock.go
+++ b/lib/go/test/tests/protocol_mock.go
@@ -60,52 +60,52 @@
 	return _mr.mock.ctrl.RecordCall(_mr.mock, "Flush")
 }
 
-func (_m *MockTProtocol) ReadBinary() ([]byte, error) {
-	ret := _m.ctrl.Call(_m, "ReadBinary")
+func (_m *MockTProtocol) ReadBinary(ctx context.Context) ([]byte, error) {
+	ret := _m.ctrl.Call(_m, "ReadBinary", ctx)
 	ret0, _ := ret[0].([]byte)
 	ret1, _ := ret[1].(error)
 	return ret0, ret1
 }
 
-func (_mr *_MockTProtocolRecorder) ReadBinary() *gomock.Call {
-	return _mr.mock.ctrl.RecordCall(_mr.mock, "ReadBinary")
+func (_mr *_MockTProtocolRecorder) ReadBinary(ctx context.Context) *gomock.Call {
+	return _mr.mock.ctrl.RecordCall(_mr.mock, "ReadBinary", ctx)
 }
 
-func (_m *MockTProtocol) ReadBool() (bool, error) {
-	ret := _m.ctrl.Call(_m, "ReadBool")
+func (_m *MockTProtocol) ReadBool(ctx context.Context) (bool, error) {
+	ret := _m.ctrl.Call(_m, "ReadBool", ctx)
 	ret0, _ := ret[0].(bool)
 	ret1, _ := ret[1].(error)
 	return ret0, ret1
 }
 
-func (_mr *_MockTProtocolRecorder) ReadBool() *gomock.Call {
-	return _mr.mock.ctrl.RecordCall(_mr.mock, "ReadBool")
+func (_mr *_MockTProtocolRecorder) ReadBool(ctx context.Context) *gomock.Call {
+	return _mr.mock.ctrl.RecordCall(_mr.mock, "ReadBool", ctx)
 }
 
-func (_m *MockTProtocol) ReadByte() (int8, error) {
-	ret := _m.ctrl.Call(_m, "ReadByte")
+func (_m *MockTProtocol) ReadByte(ctx context.Context) (int8, error) {
+	ret := _m.ctrl.Call(_m, "ReadByte", ctx)
 	ret0, _ := ret[0].(int8)
 	ret1, _ := ret[1].(error)
 	return ret0, ret1
 }
 
-func (_mr *_MockTProtocolRecorder) ReadByte() *gomock.Call {
-	return _mr.mock.ctrl.RecordCall(_mr.mock, "ReadByte")
+func (_mr *_MockTProtocolRecorder) ReadByte(ctx context.Context) *gomock.Call {
+	return _mr.mock.ctrl.RecordCall(_mr.mock, "ReadByte", ctx)
 }
 
-func (_m *MockTProtocol) ReadDouble() (float64, error) {
-	ret := _m.ctrl.Call(_m, "ReadDouble")
+func (_m *MockTProtocol) ReadDouble(ctx context.Context) (float64, error) {
+	ret := _m.ctrl.Call(_m, "ReadDouble", ctx)
 	ret0, _ := ret[0].(float64)
 	ret1, _ := ret[1].(error)
 	return ret0, ret1
 }
 
-func (_mr *_MockTProtocolRecorder) ReadDouble() *gomock.Call {
-	return _mr.mock.ctrl.RecordCall(_mr.mock, "ReadDouble")
+func (_mr *_MockTProtocolRecorder) ReadDouble(ctx context.Context) *gomock.Call {
+	return _mr.mock.ctrl.RecordCall(_mr.mock, "ReadDouble", ctx)
 }
 
-func (_m *MockTProtocol) ReadFieldBegin() (string, thrift.TType, int16, error) {
-	ret := _m.ctrl.Call(_m, "ReadFieldBegin")
+func (_m *MockTProtocol) ReadFieldBegin(ctx context.Context) (string, thrift.TType, int16, error) {
+	ret := _m.ctrl.Call(_m, "ReadFieldBegin", ctx)
 	ret0, _ := ret[0].(string)
 	ret1, _ := ret[1].(thrift.TType)
 	ret2, _ := ret[2].(int16)
@@ -113,77 +113,77 @@
 	return ret0, ret1, ret2, ret3
 }
 
-func (_mr *_MockTProtocolRecorder) ReadFieldBegin() *gomock.Call {
-	return _mr.mock.ctrl.RecordCall(_mr.mock, "ReadFieldBegin")
+func (_mr *_MockTProtocolRecorder) ReadFieldBegin(ctx context.Context) *gomock.Call {
+	return _mr.mock.ctrl.RecordCall(_mr.mock, "ReadFieldBegin", ctx)
 }
 
-func (_m *MockTProtocol) ReadFieldEnd() error {
-	ret := _m.ctrl.Call(_m, "ReadFieldEnd")
+func (_m *MockTProtocol) ReadFieldEnd(ctx context.Context) error {
+	ret := _m.ctrl.Call(_m, "ReadFieldEnd", ctx)
 	ret0, _ := ret[0].(error)
 	return ret0
 }
 
-func (_mr *_MockTProtocolRecorder) ReadFieldEnd() *gomock.Call {
-	return _mr.mock.ctrl.RecordCall(_mr.mock, "ReadFieldEnd")
+func (_mr *_MockTProtocolRecorder) ReadFieldEnd(ctx context.Context) *gomock.Call {
+	return _mr.mock.ctrl.RecordCall(_mr.mock, "ReadFieldEnd", ctx)
 }
 
-func (_m *MockTProtocol) ReadI16() (int16, error) {
-	ret := _m.ctrl.Call(_m, "ReadI16")
+func (_m *MockTProtocol) ReadI16(ctx context.Context) (int16, error) {
+	ret := _m.ctrl.Call(_m, "ReadI16", ctx)
 	ret0, _ := ret[0].(int16)
 	ret1, _ := ret[1].(error)
 	return ret0, ret1
 }
 
-func (_mr *_MockTProtocolRecorder) ReadI16() *gomock.Call {
-	return _mr.mock.ctrl.RecordCall(_mr.mock, "ReadI16")
+func (_mr *_MockTProtocolRecorder) ReadI16(ctx context.Context) *gomock.Call {
+	return _mr.mock.ctrl.RecordCall(_mr.mock, "ReadI16", ctx)
 }
 
-func (_m *MockTProtocol) ReadI32() (int32, error) {
-	ret := _m.ctrl.Call(_m, "ReadI32")
+func (_m *MockTProtocol) ReadI32(ctx context.Context) (int32, error) {
+	ret := _m.ctrl.Call(_m, "ReadI32", ctx)
 	ret0, _ := ret[0].(int32)
 	ret1, _ := ret[1].(error)
 	return ret0, ret1
 }
 
-func (_mr *_MockTProtocolRecorder) ReadI32() *gomock.Call {
-	return _mr.mock.ctrl.RecordCall(_mr.mock, "ReadI32")
+func (_mr *_MockTProtocolRecorder) ReadI32(ctx context.Context) *gomock.Call {
+	return _mr.mock.ctrl.RecordCall(_mr.mock, "ReadI32", ctx)
 }
 
-func (_m *MockTProtocol) ReadI64() (int64, error) {
-	ret := _m.ctrl.Call(_m, "ReadI64")
+func (_m *MockTProtocol) ReadI64(ctx context.Context) (int64, error) {
+	ret := _m.ctrl.Call(_m, "ReadI64", ctx)
 	ret0, _ := ret[0].(int64)
 	ret1, _ := ret[1].(error)
 	return ret0, ret1
 }
 
-func (_mr *_MockTProtocolRecorder) ReadI64() *gomock.Call {
-	return _mr.mock.ctrl.RecordCall(_mr.mock, "ReadI64")
+func (_mr *_MockTProtocolRecorder) ReadI64(ctx context.Context) *gomock.Call {
+	return _mr.mock.ctrl.RecordCall(_mr.mock, "ReadI64", ctx)
 }
 
-func (_m *MockTProtocol) ReadListBegin() (thrift.TType, int, error) {
-	ret := _m.ctrl.Call(_m, "ReadListBegin")
+func (_m *MockTProtocol) ReadListBegin(ctx context.Context) (thrift.TType, int, error) {
+	ret := _m.ctrl.Call(_m, "ReadListBegin", ctx)
 	ret0, _ := ret[0].(thrift.TType)
 	ret1, _ := ret[1].(int)
 	ret2, _ := ret[2].(error)
 	return ret0, ret1, ret2
 }
 
-func (_mr *_MockTProtocolRecorder) ReadListBegin() *gomock.Call {
-	return _mr.mock.ctrl.RecordCall(_mr.mock, "ReadListBegin")
+func (_mr *_MockTProtocolRecorder) ReadListBegin(ctx context.Context) *gomock.Call {
+	return _mr.mock.ctrl.RecordCall(_mr.mock, "ReadListBegin", ctx)
 }
 
-func (_m *MockTProtocol) ReadListEnd() error {
-	ret := _m.ctrl.Call(_m, "ReadListEnd")
+func (_m *MockTProtocol) ReadListEnd(ctx context.Context) error {
+	ret := _m.ctrl.Call(_m, "ReadListEnd", ctx)
 	ret0, _ := ret[0].(error)
 	return ret0
 }
 
-func (_mr *_MockTProtocolRecorder) ReadListEnd() *gomock.Call {
-	return _mr.mock.ctrl.RecordCall(_mr.mock, "ReadListEnd")
+func (_mr *_MockTProtocolRecorder) ReadListEnd(ctx context.Context) *gomock.Call {
+	return _mr.mock.ctrl.RecordCall(_mr.mock, "ReadListEnd", ctx)
 }
 
-func (_m *MockTProtocol) ReadMapBegin() (thrift.TType, thrift.TType, int, error) {
-	ret := _m.ctrl.Call(_m, "ReadMapBegin")
+func (_m *MockTProtocol) ReadMapBegin(ctx context.Context) (thrift.TType, thrift.TType, int, error) {
+	ret := _m.ctrl.Call(_m, "ReadMapBegin", ctx)
 	ret0, _ := ret[0].(thrift.TType)
 	ret1, _ := ret[1].(thrift.TType)
 	ret2, _ := ret[2].(int)
@@ -191,22 +191,22 @@
 	return ret0, ret1, ret2, ret3
 }
 
-func (_mr *_MockTProtocolRecorder) ReadMapBegin() *gomock.Call {
-	return _mr.mock.ctrl.RecordCall(_mr.mock, "ReadMapBegin")
+func (_mr *_MockTProtocolRecorder) ReadMapBegin(ctx context.Context) *gomock.Call {
+	return _mr.mock.ctrl.RecordCall(_mr.mock, "ReadMapBegin", ctx)
 }
 
-func (_m *MockTProtocol) ReadMapEnd() error {
-	ret := _m.ctrl.Call(_m, "ReadMapEnd")
+func (_m *MockTProtocol) ReadMapEnd(ctx context.Context) error {
+	ret := _m.ctrl.Call(_m, "ReadMapEnd", ctx)
 	ret0, _ := ret[0].(error)
 	return ret0
 }
 
-func (_mr *_MockTProtocolRecorder) ReadMapEnd() *gomock.Call {
-	return _mr.mock.ctrl.RecordCall(_mr.mock, "ReadMapEnd")
+func (_mr *_MockTProtocolRecorder) ReadMapEnd(ctx context.Context) *gomock.Call {
+	return _mr.mock.ctrl.RecordCall(_mr.mock, "ReadMapEnd", ctx)
 }
 
-func (_m *MockTProtocol) ReadMessageBegin() (string, thrift.TMessageType, int32, error) {
-	ret := _m.ctrl.Call(_m, "ReadMessageBegin")
+func (_m *MockTProtocol) ReadMessageBegin(ctx context.Context) (string, thrift.TMessageType, int32, error) {
+	ret := _m.ctrl.Call(_m, "ReadMessageBegin", ctx)
 	ret0, _ := ret[0].(string)
 	ret1, _ := ret[1].(thrift.TMessageType)
 	ret2, _ := ret[2].(int32)
@@ -214,82 +214,82 @@
 	return ret0, ret1, ret2, ret3
 }
 
-func (_mr *_MockTProtocolRecorder) ReadMessageBegin() *gomock.Call {
-	return _mr.mock.ctrl.RecordCall(_mr.mock, "ReadMessageBegin")
+func (_mr *_MockTProtocolRecorder) ReadMessageBegin(ctx context.Context) *gomock.Call {
+	return _mr.mock.ctrl.RecordCall(_mr.mock, "ReadMessageBegin", ctx)
 }
 
-func (_m *MockTProtocol) ReadMessageEnd() error {
-	ret := _m.ctrl.Call(_m, "ReadMessageEnd")
+func (_m *MockTProtocol) ReadMessageEnd(ctx context.Context) error {
+	ret := _m.ctrl.Call(_m, "ReadMessageEnd", ctx)
 	ret0, _ := ret[0].(error)
 	return ret0
 }
 
-func (_mr *_MockTProtocolRecorder) ReadMessageEnd() *gomock.Call {
-	return _mr.mock.ctrl.RecordCall(_mr.mock, "ReadMessageEnd")
+func (_mr *_MockTProtocolRecorder) ReadMessageEnd(ctx context.Context) *gomock.Call {
+	return _mr.mock.ctrl.RecordCall(_mr.mock, "ReadMessageEnd", ctx)
 }
 
-func (_m *MockTProtocol) ReadSetBegin() (thrift.TType, int, error) {
-	ret := _m.ctrl.Call(_m, "ReadSetBegin")
+func (_m *MockTProtocol) ReadSetBegin(ctx context.Context) (thrift.TType, int, error) {
+	ret := _m.ctrl.Call(_m, "ReadSetBegin", ctx)
 	ret0, _ := ret[0].(thrift.TType)
 	ret1, _ := ret[1].(int)
 	ret2, _ := ret[2].(error)
 	return ret0, ret1, ret2
 }
 
-func (_mr *_MockTProtocolRecorder) ReadSetBegin() *gomock.Call {
-	return _mr.mock.ctrl.RecordCall(_mr.mock, "ReadSetBegin")
+func (_mr *_MockTProtocolRecorder) ReadSetBegin(ctx context.Context) *gomock.Call {
+	return _mr.mock.ctrl.RecordCall(_mr.mock, "ReadSetBegin", ctx)
 }
 
-func (_m *MockTProtocol) ReadSetEnd() error {
-	ret := _m.ctrl.Call(_m, "ReadSetEnd")
+func (_m *MockTProtocol) ReadSetEnd(ctx context.Context) error {
+	ret := _m.ctrl.Call(_m, "ReadSetEnd", ctx)
 	ret0, _ := ret[0].(error)
 	return ret0
 }
 
-func (_mr *_MockTProtocolRecorder) ReadSetEnd() *gomock.Call {
-	return _mr.mock.ctrl.RecordCall(_mr.mock, "ReadSetEnd")
+func (_mr *_MockTProtocolRecorder) ReadSetEnd(ctx context.Context) *gomock.Call {
+	return _mr.mock.ctrl.RecordCall(_mr.mock, "ReadSetEnd", ctx)
 }
 
-func (_m *MockTProtocol) ReadString() (string, error) {
-	ret := _m.ctrl.Call(_m, "ReadString")
+func (_m *MockTProtocol) ReadString(ctx context.Context) (string, error) {
+	ret := _m.ctrl.Call(_m, "ReadString", ctx)
 	ret0, _ := ret[0].(string)
 	ret1, _ := ret[1].(error)
 	return ret0, ret1
 }
 
-func (_mr *_MockTProtocolRecorder) ReadString() *gomock.Call {
-	return _mr.mock.ctrl.RecordCall(_mr.mock, "ReadString")
+func (_mr *_MockTProtocolRecorder) ReadString(ctx context.Context) *gomock.Call {
+	return _mr.mock.ctrl.RecordCall(_mr.mock, "ReadString", ctx)
 }
 
-func (_m *MockTProtocol) ReadStructBegin() (string, error) {
-	ret := _m.ctrl.Call(_m, "ReadStructBegin")
+func (_m *MockTProtocol) ReadStructBegin(ctx context.Context) (string, error) {
+	ret := _m.ctrl.Call(_m, "ReadStructBegin", ctx)
 	ret0, _ := ret[0].(string)
 	ret1, _ := ret[1].(error)
 	return ret0, ret1
 }
 
-func (_mr *_MockTProtocolRecorder) ReadStructBegin() *gomock.Call {
-	return _mr.mock.ctrl.RecordCall(_mr.mock, "ReadStructBegin")
+func (_mr *_MockTProtocolRecorder) ReadStructBegin(ctx context.Context) *gomock.Call {
+	return _mr.mock.ctrl.RecordCall(_mr.mock, "ReadStructBegin", ctx)
 }
 
-func (_m *MockTProtocol) ReadStructEnd() error {
-	ret := _m.ctrl.Call(_m, "ReadStructEnd")
+func (_m *MockTProtocol) ReadStructEnd(ctx context.Context) error {
+	ret := _m.ctrl.Call(_m, "ReadStructEnd", ctx)
 	ret0, _ := ret[0].(error)
 	return ret0
 }
 
-func (_mr *_MockTProtocolRecorder) ReadStructEnd() *gomock.Call {
-	return _mr.mock.ctrl.RecordCall(_mr.mock, "ReadStructEnd")
+func (_mr *_MockTProtocolRecorder) ReadStructEnd(ctx context.Context) *gomock.Call {
+	return _mr.mock.ctrl.RecordCall(_mr.mock, "ReadStructEnd", ctx)
 }
 
-func (_m *MockTProtocol) Skip(_param0 thrift.TType) error {
-	ret := _m.ctrl.Call(_m, "Skip", _param0)
+func (_m *MockTProtocol) Skip(ctx context.Context, _param0 thrift.TType) error {
+	ret := _m.ctrl.Call(_m, "Skip", ctx, _param0)
 	ret0, _ := ret[0].(error)
 	return ret0
 }
 
-func (_mr *_MockTProtocolRecorder) Skip(arg0 interface{}) *gomock.Call {
-	return _mr.mock.ctrl.RecordCall(_mr.mock, "Skip", arg0)
+func (_mr *_MockTProtocolRecorder) Skip(ctx context.Context, arg0 interface{}) *gomock.Call {
+	return _mr.mock.ctrl.RecordCall(_mr.mock, "Skip", ctx, arg0)
 }
 
 func (_m *MockTProtocol) Transport() thrift.TTransport {
@@ -302,212 +302,212 @@
 	return _mr.mock.ctrl.RecordCall(_mr.mock, "Transport")
 }
 
-func (_m *MockTProtocol) WriteBinary(_param0 []byte) error {
-	ret := _m.ctrl.Call(_m, "WriteBinary", _param0)
+func (_m *MockTProtocol) WriteBinary(ctx context.Context, _param0 []byte) error {
+	ret := _m.ctrl.Call(_m, "WriteBinary", ctx, _param0)
 	ret0, _ := ret[0].(error)
 	return ret0
 }
 
-func (_mr *_MockTProtocolRecorder) WriteBinary(arg0 interface{}) *gomock.Call {
-	return _mr.mock.ctrl.RecordCall(_mr.mock, "WriteBinary", arg0)
+func (_mr *_MockTProtocolRecorder) WriteBinary(ctx context.Context, arg0 interface{}) *gomock.Call {
+	return _mr.mock.ctrl.RecordCall(_mr.mock, "WriteBinary", ctx, arg0)
 }
 
-func (_m *MockTProtocol) WriteBool(_param0 bool) error {
-	ret := _m.ctrl.Call(_m, "WriteBool", _param0)
+func (_m *MockTProtocol) WriteBool(ctx context.Context, _param0 bool) error {
+	ret := _m.ctrl.Call(_m, "WriteBool", ctx, _param0)
 	ret0, _ := ret[0].(error)
 	return ret0
 }
 
-func (_mr *_MockTProtocolRecorder) WriteBool(arg0 interface{}) *gomock.Call {
-	return _mr.mock.ctrl.RecordCall(_mr.mock, "WriteBool", arg0)
+func (_mr *_MockTProtocolRecorder) WriteBool(ctx context.Context, arg0 interface{}) *gomock.Call {
+	return _mr.mock.ctrl.RecordCall(_mr.mock, "WriteBool", ctx, arg0)
 }
 
-func (_m *MockTProtocol) WriteByte(_param0 int8) error {
-	ret := _m.ctrl.Call(_m, "WriteByte", _param0)
+func (_m *MockTProtocol) WriteByte(ctx context.Context, _param0 int8) error {
+	ret := _m.ctrl.Call(_m, "WriteByte", ctx, _param0)
 	ret0, _ := ret[0].(error)
 	return ret0
 }
 
-func (_mr *_MockTProtocolRecorder) WriteByte(arg0 interface{}) *gomock.Call {
-	return _mr.mock.ctrl.RecordCall(_mr.mock, "WriteByte", arg0)
+func (_mr *_MockTProtocolRecorder) WriteByte(ctx context.Context, arg0 interface{}) *gomock.Call {
+	return _mr.mock.ctrl.RecordCall(_mr.mock, "WriteByte", ctx, arg0)
 }
 
-func (_m *MockTProtocol) WriteDouble(_param0 float64) error {
-	ret := _m.ctrl.Call(_m, "WriteDouble", _param0)
+func (_m *MockTProtocol) WriteDouble(ctx context.Context, _param0 float64) error {
+	ret := _m.ctrl.Call(_m, "WriteDouble", ctx, _param0)
 	ret0, _ := ret[0].(error)
 	return ret0
 }
 
-func (_mr *_MockTProtocolRecorder) WriteDouble(arg0 interface{}) *gomock.Call {
-	return _mr.mock.ctrl.RecordCall(_mr.mock, "WriteDouble", arg0)
+func (_mr *_MockTProtocolRecorder) WriteDouble(ctx context.Context, arg0 interface{}) *gomock.Call {
+	return _mr.mock.ctrl.RecordCall(_mr.mock, "WriteDouble", ctx, arg0)
 }
 
-func (_m *MockTProtocol) WriteFieldBegin(_param0 string, _param1 thrift.TType, _param2 int16) error {
-	ret := _m.ctrl.Call(_m, "WriteFieldBegin", _param0, _param1, _param2)
+func (_m *MockTProtocol) WriteFieldBegin(ctx context.Context, _param0 string, _param1 thrift.TType, _param2 int16) error {
+	ret := _m.ctrl.Call(_m, "WriteFieldBegin", ctx, _param0, _param1, _param2)
 	ret0, _ := ret[0].(error)
 	return ret0
 }
 
-func (_mr *_MockTProtocolRecorder) WriteFieldBegin(arg0, arg1, arg2 interface{}) *gomock.Call {
-	return _mr.mock.ctrl.RecordCall(_mr.mock, "WriteFieldBegin", arg0, arg1, arg2)
+func (_mr *_MockTProtocolRecorder) WriteFieldBegin(ctx context.Context, arg0, arg1, arg2 interface{}) *gomock.Call {
+	return _mr.mock.ctrl.RecordCall(_mr.mock, "WriteFieldBegin", ctx, arg0, arg1, arg2)
 }
 
-func (_m *MockTProtocol) WriteFieldEnd() error {
-	ret := _m.ctrl.Call(_m, "WriteFieldEnd")
+func (_m *MockTProtocol) WriteFieldEnd(ctx context.Context) error {
+	ret := _m.ctrl.Call(_m, "WriteFieldEnd", ctx)
 	ret0, _ := ret[0].(error)
 	return ret0
 }
 
-func (_mr *_MockTProtocolRecorder) WriteFieldEnd() *gomock.Call {
-	return _mr.mock.ctrl.RecordCall(_mr.mock, "WriteFieldEnd")
+func (_mr *_MockTProtocolRecorder) WriteFieldEnd(ctx context.Context) *gomock.Call {
+	return _mr.mock.ctrl.RecordCall(_mr.mock, "WriteFieldEnd", ctx)
 }
 
-func (_m *MockTProtocol) WriteFieldStop() error {
-	ret := _m.ctrl.Call(_m, "WriteFieldStop")
+func (_m *MockTProtocol) WriteFieldStop(ctx context.Context) error {
+	ret := _m.ctrl.Call(_m, "WriteFieldStop", ctx)
 	ret0, _ := ret[0].(error)
 	return ret0
 }
 
-func (_mr *_MockTProtocolRecorder) WriteFieldStop() *gomock.Call {
-	return _mr.mock.ctrl.RecordCall(_mr.mock, "WriteFieldStop")
+func (_mr *_MockTProtocolRecorder) WriteFieldStop(ctx context.Context) *gomock.Call {
+	return _mr.mock.ctrl.RecordCall(_mr.mock, "WriteFieldStop", ctx)
 }
 
-func (_m *MockTProtocol) WriteI16(_param0 int16) error {
-	ret := _m.ctrl.Call(_m, "WriteI16", _param0)
+func (_m *MockTProtocol) WriteI16(ctx context.Context, _param0 int16) error {
+	ret := _m.ctrl.Call(_m, "WriteI16", ctx, _param0)
 	ret0, _ := ret[0].(error)
 	return ret0
 }
 
-func (_mr *_MockTProtocolRecorder) WriteI16(arg0 interface{}) *gomock.Call {
-	return _mr.mock.ctrl.RecordCall(_mr.mock, "WriteI16", arg0)
+func (_mr *_MockTProtocolRecorder) WriteI16(ctx context.Context, arg0 interface{}) *gomock.Call {
+	return _mr.mock.ctrl.RecordCall(_mr.mock, "WriteI16", ctx, arg0)
 }
 
-func (_m *MockTProtocol) WriteI32(_param0 int32) error {
-	ret := _m.ctrl.Call(_m, "WriteI32", _param0)
+func (_m *MockTProtocol) WriteI32(ctx context.Context, _param0 int32) error {
+	ret := _m.ctrl.Call(_m, "WriteI32", ctx, _param0)
 	ret0, _ := ret[0].(error)
 	return ret0
 }
 
-func (_mr *_MockTProtocolRecorder) WriteI32(arg0 interface{}) *gomock.Call {
-	return _mr.mock.ctrl.RecordCall(_mr.mock, "WriteI32", arg0)
+func (_mr *_MockTProtocolRecorder) WriteI32(ctx context.Context, arg0 interface{}) *gomock.Call {
+	return _mr.mock.ctrl.RecordCall(_mr.mock, "WriteI32", ctx, arg0)
 }
 
-func (_m *MockTProtocol) WriteI64(_param0 int64) error {
-	ret := _m.ctrl.Call(_m, "WriteI64", _param0)
+func (_m *MockTProtocol) WriteI64(ctx context.Context, _param0 int64) error {
+	ret := _m.ctrl.Call(_m, "WriteI64", ctx, _param0)
 	ret0, _ := ret[0].(error)
 	return ret0
 }
 
-func (_mr *_MockTProtocolRecorder) WriteI64(arg0 interface{}) *gomock.Call {
-	return _mr.mock.ctrl.RecordCall(_mr.mock, "WriteI64", arg0)
+func (_mr *_MockTProtocolRecorder) WriteI64(ctx context.Context, arg0 interface{}) *gomock.Call {
+	return _mr.mock.ctrl.RecordCall(_mr.mock, "WriteI64", ctx, arg0)
 }
 
-func (_m *MockTProtocol) WriteListBegin(_param0 thrift.TType, _param1 int) error {
-	ret := _m.ctrl.Call(_m, "WriteListBegin", _param0, _param1)
+func (_m *MockTProtocol) WriteListBegin(ctx context.Context, _param0 thrift.TType, _param1 int) error {
+	ret := _m.ctrl.Call(_m, "WriteListBegin", ctx, _param0, _param1)
 	ret0, _ := ret[0].(error)
 	return ret0
 }
 
-func (_mr *_MockTProtocolRecorder) WriteListBegin(arg0, arg1 interface{}) *gomock.Call {
-	return _mr.mock.ctrl.RecordCall(_mr.mock, "WriteListBegin", arg0, arg1)
+func (_mr *_MockTProtocolRecorder) WriteListBegin(ctx context.Context, arg0, arg1 interface{}) *gomock.Call {
+	return _mr.mock.ctrl.RecordCall(_mr.mock, "WriteListBegin", ctx, arg0, arg1)
 }
 
-func (_m *MockTProtocol) WriteListEnd() error {
-	ret := _m.ctrl.Call(_m, "WriteListEnd")
+func (_m *MockTProtocol) WriteListEnd(ctx context.Context) error {
+	ret := _m.ctrl.Call(_m, "WriteListEnd", ctx)
 	ret0, _ := ret[0].(error)
 	return ret0
 }
 
-func (_mr *_MockTProtocolRecorder) WriteListEnd() *gomock.Call {
-	return _mr.mock.ctrl.RecordCall(_mr.mock, "WriteListEnd")
+func (_mr *_MockTProtocolRecorder) WriteListEnd(ctx context.Context) *gomock.Call {
+	return _mr.mock.ctrl.RecordCall(_mr.mock, "WriteListEnd", ctx)
 }
 
-func (_m *MockTProtocol) WriteMapBegin(_param0 thrift.TType, _param1 thrift.TType, _param2 int) error {
-	ret := _m.ctrl.Call(_m, "WriteMapBegin", _param0, _param1, _param2)
+func (_m *MockTProtocol) WriteMapBegin(ctx context.Context, _param0 thrift.TType, _param1 thrift.TType, _param2 int) error {
+	ret := _m.ctrl.Call(_m, "WriteMapBegin", ctx, _param0, _param1, _param2)
 	ret0, _ := ret[0].(error)
 	return ret0
 }
 
-func (_mr *_MockTProtocolRecorder) WriteMapBegin(arg0, arg1, arg2 interface{}) *gomock.Call {
-	return _mr.mock.ctrl.RecordCall(_mr.mock, "WriteMapBegin", arg0, arg1, arg2)
+func (_mr *_MockTProtocolRecorder) WriteMapBegin(ctx context.Context, arg0, arg1, arg2 interface{}) *gomock.Call {
+	return _mr.mock.ctrl.RecordCall(_mr.mock, "WriteMapBegin", ctx, arg0, arg1, arg2)
 }
 
-func (_m *MockTProtocol) WriteMapEnd() error {
-	ret := _m.ctrl.Call(_m, "WriteMapEnd")
+func (_m *MockTProtocol) WriteMapEnd(ctx context.Context) error {
+	ret := _m.ctrl.Call(_m, "WriteMapEnd", ctx)
 	ret0, _ := ret[0].(error)
 	return ret0
 }
 
-func (_mr *_MockTProtocolRecorder) WriteMapEnd() *gomock.Call {
-	return _mr.mock.ctrl.RecordCall(_mr.mock, "WriteMapEnd")
+func (_mr *_MockTProtocolRecorder) WriteMapEnd(ctx context.Context) *gomock.Call {
+	return _mr.mock.ctrl.RecordCall(_mr.mock, "WriteMapEnd", ctx)
 }
 
-func (_m *MockTProtocol) WriteMessageBegin(_param0 string, _param1 thrift.TMessageType, _param2 int32) error {
-	ret := _m.ctrl.Call(_m, "WriteMessageBegin", _param0, _param1, _param2)
+func (_m *MockTProtocol) WriteMessageBegin(ctx context.Context, _param0 string, _param1 thrift.TMessageType, _param2 int32) error {
+	ret := _m.ctrl.Call(_m, "WriteMessageBegin", ctx, _param0, _param1, _param2)
 	ret0, _ := ret[0].(error)
 	return ret0
 }
 
-func (_mr *_MockTProtocolRecorder) WriteMessageBegin(arg0, arg1, arg2 interface{}) *gomock.Call {
-	return _mr.mock.ctrl.RecordCall(_mr.mock, "WriteMessageBegin", arg0, arg1, arg2)
+func (_mr *_MockTProtocolRecorder) WriteMessageBegin(ctx context.Context, arg0, arg1, arg2 interface{}) *gomock.Call {
+	return _mr.mock.ctrl.RecordCall(_mr.mock, "WriteMessageBegin", ctx, arg0, arg1, arg2)
 }
 
-func (_m *MockTProtocol) WriteMessageEnd() error {
-	ret := _m.ctrl.Call(_m, "WriteMessageEnd")
+func (_m *MockTProtocol) WriteMessageEnd(ctx context.Context) error {
+	ret := _m.ctrl.Call(_m, "WriteMessageEnd", ctx)
 	ret0, _ := ret[0].(error)
 	return ret0
 }
 
-func (_mr *_MockTProtocolRecorder) WriteMessageEnd() *gomock.Call {
-	return _mr.mock.ctrl.RecordCall(_mr.mock, "WriteMessageEnd")
+func (_mr *_MockTProtocolRecorder) WriteMessageEnd(ctx context.Context) *gomock.Call {
+	return _mr.mock.ctrl.RecordCall(_mr.mock, "WriteMessageEnd", ctx)
 }
 
-func (_m *MockTProtocol) WriteSetBegin(_param0 thrift.TType, _param1 int) error {
-	ret := _m.ctrl.Call(_m, "WriteSetBegin", _param0, _param1)
+func (_m *MockTProtocol) WriteSetBegin(ctx context.Context, _param0 thrift.TType, _param1 int) error {
+	ret := _m.ctrl.Call(_m, "WriteSetBegin", ctx, _param0, _param1)
 	ret0, _ := ret[0].(error)
 	return ret0
 }
 
-func (_mr *_MockTProtocolRecorder) WriteSetBegin(arg0, arg1 interface{}) *gomock.Call {
-	return _mr.mock.ctrl.RecordCall(_mr.mock, "WriteSetBegin", arg0, arg1)
+func (_mr *_MockTProtocolRecorder) WriteSetBegin(ctx context.Context, arg0, arg1 interface{}) *gomock.Call {
+	return _mr.mock.ctrl.RecordCall(_mr.mock, "WriteSetBegin", ctx, arg0, arg1)
 }
 
-func (_m *MockTProtocol) WriteSetEnd() error {
-	ret := _m.ctrl.Call(_m, "WriteSetEnd")
+func (_m *MockTProtocol) WriteSetEnd(ctx context.Context) error {
+	ret := _m.ctrl.Call(_m, "WriteSetEnd", ctx)
 	ret0, _ := ret[0].(error)
 	return ret0
 }
 
-func (_mr *_MockTProtocolRecorder) WriteSetEnd() *gomock.Call {
-	return _mr.mock.ctrl.RecordCall(_mr.mock, "WriteSetEnd")
+func (_mr *_MockTProtocolRecorder) WriteSetEnd(ctx context.Context) *gomock.Call {
+	return _mr.mock.ctrl.RecordCall(_mr.mock, "WriteSetEnd", ctx)
 }
 
-func (_m *MockTProtocol) WriteString(_param0 string) error {
-	ret := _m.ctrl.Call(_m, "WriteString", _param0)
+func (_m *MockTProtocol) WriteString(ctx context.Context, _param0 string) error {
+	ret := _m.ctrl.Call(_m, "WriteString", ctx, _param0)
 	ret0, _ := ret[0].(error)
 	return ret0
 }
 
-func (_mr *_MockTProtocolRecorder) WriteString(arg0 interface{}) *gomock.Call {
-	return _mr.mock.ctrl.RecordCall(_mr.mock, "WriteString", arg0)
+func (_mr *_MockTProtocolRecorder) WriteString(ctx context.Context, arg0 interface{}) *gomock.Call {
+	return _mr.mock.ctrl.RecordCall(_mr.mock, "WriteString", ctx, arg0)
 }
 
-func (_m *MockTProtocol) WriteStructBegin(_param0 string) error {
-	ret := _m.ctrl.Call(_m, "WriteStructBegin", _param0)
+func (_m *MockTProtocol) WriteStructBegin(ctx context.Context, _param0 string) error {
+	ret := _m.ctrl.Call(_m, "WriteStructBegin", ctx, _param0)
 	ret0, _ := ret[0].(error)
 	return ret0
 }
 
-func (_mr *_MockTProtocolRecorder) WriteStructBegin(arg0 interface{}) *gomock.Call {
-	return _mr.mock.ctrl.RecordCall(_mr.mock, "WriteStructBegin", arg0)
+func (_mr *_MockTProtocolRecorder) WriteStructBegin(ctx context.Context, arg0 interface{}) *gomock.Call {
+	return _mr.mock.ctrl.RecordCall(_mr.mock, "WriteStructBegin", ctx, arg0)
 }
 
-func (_m *MockTProtocol) WriteStructEnd() error {
-	ret := _m.ctrl.Call(_m, "WriteStructEnd")
+func (_m *MockTProtocol) WriteStructEnd(ctx context.Context) error {
+	ret := _m.ctrl.Call(_m, "WriteStructEnd", ctx)
 	ret0, _ := ret[0].(error)
 	return ret0
 }
 
-func (_mr *_MockTProtocolRecorder) WriteStructEnd() *gomock.Call {
-	return _mr.mock.ctrl.RecordCall(_mr.mock, "WriteStructEnd")
+func (_mr *_MockTProtocolRecorder) WriteStructEnd(ctx context.Context) *gomock.Call {
+	return _mr.mock.ctrl.RecordCall(_mr.mock, "WriteStructEnd", ctx)
 }
diff --git a/lib/go/test/tests/required_fields_test.go b/lib/go/test/tests/required_fields_test.go
index 3fa414a..06e8560 100644
--- a/lib/go/test/tests/required_fields_test.go
+++ b/lib/go/test/tests/required_fields_test.go
@@ -37,7 +37,7 @@
 	}
 
 	d := thrift.NewTDeserializer()
-	err = d.Read(&requiredfieldtest.RequiredField{}, sourceData)
+	err = d.Read(context.Background(), &requiredfieldtest.RequiredField{}, sourceData)
 	if err != nil {
 		t.Fatalf("Did not expect an error when trying to deserialize the requiredfieldtest.RequiredField: %v", err)
 	}
@@ -53,7 +53,7 @@
 
 	// attempt to deserialize into a different type (which should fail)
 	d := thrift.NewTDeserializer()
-	err = d.Read(&requiredfieldtest.RequiredField{}, sourceData)
+	err = d.Read(context.Background(), &requiredfieldtest.RequiredField{}, sourceData)
 	if err == nil {
 		t.Fatal("Expected an error when trying to deserialize an object which is missing a required field")
 	}
@@ -66,12 +66,12 @@
 
 	// None of required fields are set
 	gomock.InOrder(
-		protocol.EXPECT().ReadStructBegin().Return("StructC", nil),
-		protocol.EXPECT().ReadFieldBegin().Return("_", thrift.TType(thrift.STOP), int16(1), nil),
-		protocol.EXPECT().ReadStructEnd().Return(nil),
+		protocol.EXPECT().ReadStructBegin(context.Background()).Return("StructC", nil),
+		protocol.EXPECT().ReadFieldBegin(context.Background()).Return("_", thrift.TType(thrift.STOP), int16(1), nil),
+		protocol.EXPECT().ReadStructEnd(context.Background()).Return(nil),
 	)
 
-	err := testStruct.Read(protocol)
+	err := testStruct.Read(context.Background(), protocol)
 	mockCtrl.Finish()
 	mockCtrl = gomock.NewController(t)
 	if err == nil {
@@ -87,15 +87,15 @@
 
 	// One of the required fields is set
 	gomock.InOrder(
-		protocol.EXPECT().ReadStructBegin().Return("StructC", nil),
-		protocol.EXPECT().ReadFieldBegin().Return("I", thrift.TType(thrift.I32), int16(2), nil),
-		protocol.EXPECT().ReadI32().Return(int32(1), nil),
-		protocol.EXPECT().ReadFieldEnd().Return(nil),
-		protocol.EXPECT().ReadFieldBegin().Return("_", thrift.TType(thrift.STOP), int16(1), nil),
-		protocol.EXPECT().ReadStructEnd().Return(nil),
+		protocol.EXPECT().ReadStructBegin(context.Background()).Return("StructC", nil),
+		protocol.EXPECT().ReadFieldBegin(context.Background()).Return("I", thrift.TType(thrift.I32), int16(2), nil),
+		protocol.EXPECT().ReadI32(context.Background()).Return(int32(1), nil),
+		protocol.EXPECT().ReadFieldEnd(context.Background()).Return(nil),
+		protocol.EXPECT().ReadFieldBegin(context.Background()).Return("_", thrift.TType(thrift.STOP), int16(1), nil),
+		protocol.EXPECT().ReadStructEnd(context.Background()).Return(nil),
 	)
 
-	err = testStruct.Read(protocol)
+	err = testStruct.Read(context.Background(), protocol)
 	mockCtrl.Finish()
 	mockCtrl = gomock.NewController(t)
 	if err == nil {
@@ -111,18 +111,18 @@
 
 	// Both of the required fields are set
 	gomock.InOrder(
-		protocol.EXPECT().ReadStructBegin().Return("StructC", nil),
-		protocol.EXPECT().ReadFieldBegin().Return("i", thrift.TType(thrift.I32), int16(2), nil),
-		protocol.EXPECT().ReadI32().Return(int32(1), nil),
-		protocol.EXPECT().ReadFieldEnd().Return(nil),
-		protocol.EXPECT().ReadFieldBegin().Return("s2", thrift.TType(thrift.STRING), int16(4), nil),
-		protocol.EXPECT().ReadString().Return("test", nil),
-		protocol.EXPECT().ReadFieldEnd().Return(nil),
-		protocol.EXPECT().ReadFieldBegin().Return("_", thrift.TType(thrift.STOP), int16(1), nil),
-		protocol.EXPECT().ReadStructEnd().Return(nil),
+		protocol.EXPECT().ReadStructBegin(context.Background()).Return("StructC", nil),
+		protocol.EXPECT().ReadFieldBegin(context.Background()).Return("i", thrift.TType(thrift.I32), int16(2), nil),
+		protocol.EXPECT().ReadI32(context.Background()).Return(int32(1), nil),
+		protocol.EXPECT().ReadFieldEnd(context.Background()).Return(nil),
+		protocol.EXPECT().ReadFieldBegin(context.Background()).Return("s2", thrift.TType(thrift.STRING), int16(4), nil),
+		protocol.EXPECT().ReadString(context.Background()).Return("test", nil),
+		protocol.EXPECT().ReadFieldEnd(context.Background()).Return(nil),
+		protocol.EXPECT().ReadFieldBegin(context.Background()).Return("_", thrift.TType(thrift.STOP), int16(1), nil),
+		protocol.EXPECT().ReadStructEnd(context.Background()).Return(nil),
 	)
 
-	err = testStruct.Read(protocol)
+	err = testStruct.Read(context.Background(), protocol)
 	mockCtrl.Finish()
 	if err != nil {
 		t.Fatal("Expected read to succeed")