THRIFT-4914: Add TResponseHelper

This is the third part of THRIFT-4914, which handles the server writing
part in the response (server -> client direction).

Define a new type, TResponseHelper, which only contains THeader related
functions for now, but can be extended for other functions in the
future.

In TSimpleServer, inject a TResponseHelper into the context object
passed into the handler functions. Handler function code could retrieve
the injected TResponseHelper to set headers to be written to the client.

Client: go

This closes #1923.

lib/go/thrift/response_helper.go

@@ -0,0 +1,94 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 thrift
+
+import (
+	"context"
+)
+
+// See https://godoc.org/context#WithValue on why do we need the unexported typedefs.
+type responseHelperKey struct{}
+
+// TResponseHelper defines a object with a set of helper functions that can be
+// retrieved from the context object passed into server handler functions.
+//
+// Use GetResponseHelper to retrieve the injected TResponseHelper implementation
+// from the context object.
+//
+// The zero value of TResponseHelper is valid with all helper functions being
+// no-op.
+type TResponseHelper struct {
+	// THeader related functions
+	*THeaderResponseHelper
+}
+
+// THeaderResponseHelper defines THeader related TResponseHelper functions.
+//
+// The zero value of *THeaderResponseHelper is valid with all helper functions
+// being no-op.
+type THeaderResponseHelper struct {
+	proto *THeaderProtocol
+}
+
+// NewTHeaderResponseHelper creates a new THeaderResponseHelper from the
+// underlying TProtocol.
+func NewTHeaderResponseHelper(proto TProtocol) *THeaderResponseHelper {
+	if hp, ok := proto.(*THeaderProtocol); ok {
+		return &THeaderResponseHelper{
+			proto: hp,
+		}
+	}
+	return nil
+}
+
+// SetHeader sets a response header.
+//
+// It's no-op if the underlying protocol/transport does not support THeader.
+func (h *THeaderResponseHelper) SetHeader(key, value string) {
+	if h != nil && h.proto != nil {
+		h.proto.SetWriteHeader(key, value)
+	}
+}
+
+// ClearHeaders clears all the response headers previously set.
+//
+// It's no-op if the underlying protocol/transport does not support THeader.
+func (h *THeaderResponseHelper) ClearHeaders() {
+	if h != nil && h.proto != nil {
+		h.proto.ClearWriteHeaders()
+	}
+}
+
+// GetResponseHelper retrieves the TResponseHelper implementation injected into
+// the context object.
+//
+// If no helper was found in the context object, a nop helper with ok == false
+// will be returned.
+func GetResponseHelper(ctx context.Context) (helper TResponseHelper, ok bool) {
+	if v := ctx.Value(responseHelperKey{}); v != nil {
+		helper, ok = v.(TResponseHelper)
+	}
+	return
+}
+
+// SetResponseHelper injects TResponseHelper into the context object.
+func SetResponseHelper(ctx context.Context, helper TResponseHelper) context.Context {
+	return context.WithValue(ctx, responseHelperKey{}, helper)
+}

lib/go/thrift/response_helper_test.go

@@ -0,0 +1,137 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 thrift
+
+import (
+	"context"
+	"testing"
+)
+
+func TestResponseHelperContext(t *testing.T) {
+	ctx := context.Background()
+
+	t.Run(
+		"empty-noop",
+		func(t *testing.T) {
+			helper, ok := GetResponseHelper(ctx)
+			if ok {
+				t.Error("GetResponseHelper expected ok == false")
+			}
+			// Just make sure those function calls does not panic
+			helper.SetHeader("foo", "bar")
+			helper.ClearHeaders()
+		},
+	)
+
+	t.Run(
+		"set-get",
+		func(t *testing.T) {
+			trans := NewTHeaderTransport(NewTMemoryBuffer())
+			proto := NewTHeaderProtocol(trans)
+			ctx = SetResponseHelper(
+				ctx,
+				TResponseHelper{
+					THeaderResponseHelper: NewTHeaderResponseHelper(proto),
+				},
+			)
+			helper, ok := GetResponseHelper(ctx)
+			if !ok {
+				t.Error("GetResponseHelper expected ok == true")
+			}
+			if helper.THeaderResponseHelper == nil {
+				t.Error("GetResponseHelper expected THeaderResponseHelper to be non-nil")
+			}
+		},
+	)
+}
+
+func TestHeaderHelper(t *testing.T) {
+	t.Run(
+		"THeaderProtocol",
+		func(t *testing.T) {
+			trans := NewTHeaderTransport(NewTMemoryBuffer())
+			proto := NewTHeaderProtocol(trans)
+			helper := NewTHeaderResponseHelper(proto)
+
+			const (
+				key   = "key"
+				value = "value"
+			)
+			helper.SetHeader(key, value)
+			if len(trans.writeHeaders) != 1 {
+				t.Errorf(
+					"Expected THeaderTransport.writeHeaders to be with size of 1, got %+v",
+					trans.writeHeaders,
+				)
+			}
+			actual := trans.writeHeaders[key]
+			if actual != value {
+				t.Errorf(
+					"Expected THeaderTransport.writeHeaders to have %q:%q, got %+v",
+					key,
+					value,
+					trans.writeHeaders,
+				)
+			}
+			helper.ClearHeaders()
+			if len(trans.writeHeaders) != 0 {
+				t.Errorf(
+					"Expected THeaderTransport.writeHeaders to be empty after ClearHeaders call, got %+v",
+					trans.writeHeaders,
+				)
+			}
+		},
+	)
+
+	t.Run(
+		"other-protocol",
+		func(t *testing.T) {
+			trans := NewTMemoryBuffer()
+			proto := NewTCompactProtocol(trans)
+			helper := NewTHeaderResponseHelper(proto)
+
+			// We only need to make sure that functions in helper
+			// don't panic here.
+			helper.SetHeader("foo", "bar")
+			helper.ClearHeaders()
+		},
+	)
+
+	t.Run(
+		"zero-value",
+		func(t *testing.T) {
+			var helper *THeaderResponseHelper
+
+			// We only need to make sure that functions in helper
+			// don't panic here.
+			helper.SetHeader("foo", "bar")
+			helper.ClearHeaders()
+		},
+	)
+}
+
+func TestTResponseHelperZeroValue(t *testing.T) {
+	var helper THeaderResponseHelper
+
+	// We only need to make sure that functions in helper
+	// don't panic here.
+	helper.SetHeader("foo", "bar")
+	helper.ClearHeaders()
+}

lib/go/thrift/simple_server.go

@@ -272,7 +272,12 @@ func (p *TSimpleServer) processRequests(client TTransport) (err error) {
 			return nil
 		}
 
-		ctx := defaultCtx
+		ctx := SetResponseHelper(
+			defaultCtx,
+			TResponseHelper{
+				THeaderResponseHelper: NewTHeaderResponseHelper(outputProtocol),
+			},
+		)
 		if headerProtocol != nil {
 			// We need to call ReadFrame here, otherwise we won't
 			// get any headers on the AddReadTHeaderToContext call.