THRIFT-4914: Send context THeaders for client writes
Client: go
This is the second part of THRIFT-4914, which handles the client writing
part in the requests (client -> server direction).
In TStandardClient, when the context has write headers set, and the
protocol is THeaderProtocol, automatically extract all headers from the
context object and set to THeaderProtocol to send over the wire.
Client code can set headers into the context object by using the helper
functions in header_context.go.
Note that we have separated keys for read and write header key list, so
that for code that's both a server and a client (example: a server that
calls other upstream thrift servers), they don't automatically forward
all headers to their upstream servers, and need to explicitly set which
headers to forward.
In order to make auto forwarding easier, also add SetForwardHeaders
function to TSimpleServer, which will help the users to auto forward
selected headers.
This closes #1845.
diff --git a/lib/go/thrift/client.go b/lib/go/thrift/client.go
index 28791cc..b073a95 100644
--- a/lib/go/thrift/client.go
+++ b/lib/go/thrift/client.go
@@ -24,6 +24,16 @@
}
func (p *TStandardClient) Send(ctx context.Context, oprot TProtocol, seqId int32, method string, args TStruct) error {
+ // Set headers from context object on THeaderProtocol
+ if headerProt, ok := oprot.(*THeaderProtocol); ok {
+ headerProt.ClearWriteHeaders()
+ for _, key := range GetWriteHeaderList(ctx) {
+ if value, ok := GetHeader(ctx, key); ok {
+ headerProt.SetWriteHeader(key, value)
+ }
+ }
+ }
+
if err := oprot.WriteMessageBegin(method, CALL, seqId); err != nil {
return err
}
diff --git a/lib/go/thrift/header_context.go b/lib/go/thrift/header_context.go
index 5d9104b..21e880d 100644
--- a/lib/go/thrift/header_context.go
+++ b/lib/go/thrift/header_context.go
@@ -32,6 +32,7 @@
// Values for headerKeyList.
const (
headerKeyListRead headerKeyList = iota
+ headerKeyListWrite
)
// SetHeader sets a header in the context.
@@ -70,6 +71,25 @@
return nil
}
+// SetWriteHeaderList sets the key list of THeaders to write in the context.
+func SetWriteHeaderList(ctx context.Context, keys []string) context.Context {
+ return context.WithValue(
+ ctx,
+ headerKeyListWrite,
+ keys,
+ )
+}
+
+// GetWriteHeaderList returns the key list of THeaders to write from the context.
+func GetWriteHeaderList(ctx context.Context) []string {
+ if v := ctx.Value(headerKeyListWrite); v != nil {
+ if value, ok := v.([]string); ok {
+ return value
+ }
+ }
+ return nil
+}
+
// AddReadTHeaderToContext adds the whole THeader headers into context.
func AddReadTHeaderToContext(ctx context.Context, headers THeaderMap) context.Context {
keys := make([]string, 0, len(headers))
diff --git a/lib/go/thrift/header_context_test.go b/lib/go/thrift/header_context_test.go
index 33ac4ec..a1ea2d0 100644
--- a/lib/go/thrift/header_context_test.go
+++ b/lib/go/thrift/header_context_test.go
@@ -70,7 +70,7 @@
)
}
-func TestKeyList(t *testing.T) {
+func TestReadKeyList(t *testing.T) {
headers := THeaderMap{
"key1": "value1",
"key2": "value2",
@@ -94,4 +94,35 @@
if !reflect.DeepEqual(headers, got) {
t.Errorf("Expected header map %+v, got %+v", headers, got)
}
+
+ writtenKeys := GetWriteHeaderList(ctx)
+ if len(writtenKeys) > 0 {
+ t.Errorf(
+ "Expected empty GetWriteHeaderList() result, got %+v",
+ writtenKeys,
+ )
+ }
+}
+
+func TestWriteKeyList(t *testing.T) {
+ keys := []string{
+ "key1",
+ "key2",
+ }
+ ctx := context.Background()
+
+ ctx = SetWriteHeaderList(ctx, keys)
+ got := GetWriteHeaderList(ctx)
+
+ if !reflect.DeepEqual(keys, got) {
+ t.Errorf("Expected header keys %+v, got %+v", keys, got)
+ }
+
+ readKeys := GetReadHeaderList(ctx)
+ if len(readKeys) > 0 {
+ t.Errorf(
+ "Expected empty GetReadHeaderList() result, got %+v",
+ readKeys,
+ )
+ }
}
diff --git a/lib/go/thrift/simple_server.go b/lib/go/thrift/simple_server.go
index 9155cfb..f8efbed 100644
--- a/lib/go/thrift/simple_server.go
+++ b/lib/go/thrift/simple_server.go
@@ -42,6 +42,9 @@
outputTransportFactory TTransportFactory
inputProtocolFactory TProtocolFactory
outputProtocolFactory TProtocolFactory
+
+ // Headers to auto forward in THeaderProtocol
+ forwardHeaders []string
}
func NewTSimpleServer2(processor TProcessor, serverTransport TServerTransport) *TSimpleServer {
@@ -125,6 +128,26 @@
return p.serverTransport.Listen()
}
+// SetForwardHeaders sets the list of header keys that will be auto forwarded
+// while using THeaderProtocol.
+//
+// "forward" means that when the server is also a client to other upstream
+// thrift servers, the context object user gets in the processor functions will
+// have both read and write headers set, with write headers being forwarded.
+// Users can always override the write headers by calling SetWriteHeaderList
+// before calling thrift client functions.
+func (p *TSimpleServer) SetForwardHeaders(headers []string) {
+ size := len(headers)
+ if size == 0 {
+ p.forwardHeaders = nil
+ return
+ }
+
+ keys := make([]string, size)
+ copy(keys, headers)
+ p.forwardHeaders = keys
+}
+
func (p *TSimpleServer) innerAccept() (int32, error) {
client, err := p.serverTransport.Accept()
p.mu.Lock()
@@ -235,6 +258,7 @@
return err
}
ctx = AddReadTHeaderToContext(defaultCtx, headerProtocol.GetReadHeaders())
+ ctx = SetWriteHeaderList(ctx, p.forwardHeaders)
}
ok, err := processor.Process(ctx, inputProtocol, outputProtocol)