From fd3cb964032e0cfd6d3f992bdba46dceab21f1d0 Mon Sep 17 00:00:00 2001 From: Dorian Jaminais-Grellier Date: Fri, 5 Jan 2024 13:34:37 -0800 Subject: [PATCH 1/6] Add config for the header use by the general ingress --- thriftbp/client_pool.go | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/thriftbp/client_pool.go b/thriftbp/client_pool.go index 1f307d921..1cf808543 100644 --- a/thriftbp/client_pool.go +++ b/thriftbp/client_pool.go @@ -211,6 +211,11 @@ type ClientPoolConfig struct { // // Optional. If this is empty, no "User-Agent" header will be sent. ClientName string `yaml:"clientName"` + + // The hostname to add as a "thrift-hostname" header. + // + // Optional. If empty, not "thrift-hostname" header will be sent. + ThriftHostnameHeader string `yaml:"thriftHostnameHeader"` } // Validate checks ClientPoolConfig for any missing or erroneous values. @@ -429,6 +434,19 @@ func NewCustomClientPoolWithContext( return newClientPool(ctx, cfg, genAddr, protoFactory, middlewares...) } +const ThriftHostnameHeader = "thrift-hostname" + +func ThriftHostnameHeaderMiddleware(hostname string) thrift.ClientMiddleware { + return func(next thrift.TClient) thrift.TClient { + return thrift.WrappedTClient{ + Wrapped: func(ctx context.Context, method string, args, result thrift.TStruct) (thrift.ResponseMeta, error) { + ctx = AddClientHeader(ctx, ThriftHostnameHeader, hostname) + return next.Call(ctx, method, args, result) + }, + } + } +} + func newClientPool( ctx context.Context, cfg ClientPoolConfig, @@ -505,6 +523,11 @@ func newClientPool( slug: cfg.ServiceSlug, } + + if cfg.ThriftHostnameHeader != "" { + middlewares = append(middlewares, ThriftHostnameHeaderMiddleware(cfg.ThriftHostnameHeader)) + } + // finish setting up the clientPool by wrapping the inner "Call" with the // given middleware. // From 6d1a7630243436611e54c9603ef91fa5a3b6d38e Mon Sep 17 00:00:00 2001 From: Dorian Jaminais-Grellier Date: Fri, 5 Jan 2024 13:51:31 -0800 Subject: [PATCH 2/6] always add middleware but don't always add the header --- thriftbp/client_pool.go | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/thriftbp/client_pool.go b/thriftbp/client_pool.go index 1cf808543..001962ef2 100644 --- a/thriftbp/client_pool.go +++ b/thriftbp/client_pool.go @@ -440,7 +440,9 @@ func ThriftHostnameHeaderMiddleware(hostname string) thrift.ClientMiddleware { return func(next thrift.TClient) thrift.TClient { return thrift.WrappedTClient{ Wrapped: func(ctx context.Context, method string, args, result thrift.TStruct) (thrift.ResponseMeta, error) { - ctx = AddClientHeader(ctx, ThriftHostnameHeader, hostname) + if hostname != "" { + ctx = AddClientHeader(ctx, ThriftHostnameHeader, hostname) + } return next.Call(ctx, method, args, result) }, } @@ -523,10 +525,7 @@ func newClientPool( slug: cfg.ServiceSlug, } - - if cfg.ThriftHostnameHeader != "" { - middlewares = append(middlewares, ThriftHostnameHeaderMiddleware(cfg.ThriftHostnameHeader)) - } + middlewares = append(middlewares, ThriftHostnameHeaderMiddleware(cfg.ThriftHostnameHeader)) // finish setting up the clientPool by wrapping the inner "Call" with the // given middleware. From 5079a2cf58a2e1d9aa5ec1360571ed4ff73d7711 Mon Sep 17 00:00:00 2001 From: Dorian Jaminais-Grellier Date: Fri, 5 Jan 2024 13:59:33 -0800 Subject: [PATCH 3/6] capitalize the header --- thriftbp/client_pool.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/thriftbp/client_pool.go b/thriftbp/client_pool.go index 001962ef2..1039845e1 100644 --- a/thriftbp/client_pool.go +++ b/thriftbp/client_pool.go @@ -434,7 +434,7 @@ func NewCustomClientPoolWithContext( return newClientPool(ctx, cfg, genAddr, protoFactory, middlewares...) } -const ThriftHostnameHeader = "thrift-hostname" +const ThriftHostnameHeader = "Thrift-Hostname" func ThriftHostnameHeaderMiddleware(hostname string) thrift.ClientMiddleware { return func(next thrift.TClient) thrift.TClient { From 8d97cf0ba917b07885bfa6212a0c8aa6a6dd2561 Mon Sep 17 00:00:00 2001 From: Dorian Jaminais-Grellier Date: Fri, 5 Jan 2024 15:05:27 -0800 Subject: [PATCH 4/6] use lower case --- thriftbp/client_pool.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/thriftbp/client_pool.go b/thriftbp/client_pool.go index 1039845e1..001962ef2 100644 --- a/thriftbp/client_pool.go +++ b/thriftbp/client_pool.go @@ -434,7 +434,7 @@ func NewCustomClientPoolWithContext( return newClientPool(ctx, cfg, genAddr, protoFactory, middlewares...) } -const ThriftHostnameHeader = "Thrift-Hostname" +const ThriftHostnameHeader = "thrift-hostname" func ThriftHostnameHeaderMiddleware(hostname string) thrift.ClientMiddleware { return func(next thrift.TClient) thrift.TClient { From 1cb242594bb81286ff548eb8bf99c24b20c3f7b4 Mon Sep 17 00:00:00 2001 From: Dorian Jaminais-Grellier Date: Mon, 8 Jan 2024 08:58:55 -0800 Subject: [PATCH 5/6] address some of the feedback --- thriftbp/client_pool.go | 9 ++++++--- thriftbp/client_pool_test.go | 33 +++++++++++++++++++++++++++++++++ 2 files changed, 39 insertions(+), 3 deletions(-) diff --git a/thriftbp/client_pool.go b/thriftbp/client_pool.go index 001962ef2..27a66efb9 100644 --- a/thriftbp/client_pool.go +++ b/thriftbp/client_pool.go @@ -214,7 +214,7 @@ type ClientPoolConfig struct { // The hostname to add as a "thrift-hostname" header. // - // Optional. If empty, not "thrift-hostname" header will be sent. + // Optional. If empty, no "thrift-hostname" header will be sent. ThriftHostnameHeader string `yaml:"thriftHostnameHeader"` } @@ -436,7 +436,10 @@ func NewCustomClientPoolWithContext( const ThriftHostnameHeader = "thrift-hostname" -func ThriftHostnameHeaderMiddleware(hostname string) thrift.ClientMiddleware { +// thriftHostnameHeaderMiddleware adds a `thrift-hostname` header if one was +// specified in the configuration. +// This middleware is always added but will only add the header is necessary. +func thriftHostnameHeaderMiddleware(hostname string) thrift.ClientMiddleware { return func(next thrift.TClient) thrift.TClient { return thrift.WrappedTClient{ Wrapped: func(ctx context.Context, method string, args, result thrift.TStruct) (thrift.ResponseMeta, error) { @@ -525,7 +528,7 @@ func newClientPool( slug: cfg.ServiceSlug, } - middlewares = append(middlewares, ThriftHostnameHeaderMiddleware(cfg.ThriftHostnameHeader)) + middlewares = append(middlewares, thriftHostnameHeaderMiddleware(cfg.ThriftHostnameHeader)) // finish setting up the clientPool by wrapping the inner "Call" with the // given middleware. diff --git a/thriftbp/client_pool_test.go b/thriftbp/client_pool_test.go index 71e0de803..0650b95d1 100644 --- a/thriftbp/client_pool_test.go +++ b/thriftbp/client_pool_test.go @@ -190,6 +190,39 @@ func TestBehaviorWithNetworkIssues(t *testing.T) { } } +func TestThriftHostnameHeader(t *testing.T) { + ln, err := net.Listen("tcp", addr) + if err != nil { + t.Fatal(err) + } + defer ln.Close() + + cfg := thriftbp.ClientPoolConfig{ + Addr: ":9090", + EdgeContextImpl: ecinterface.Mock(), + ServiceSlug: "test", + InitialConnections: 1, + MaxConnections: 5, + ConnectTimeout: time.Millisecond * 5, + SocketTimeout: time.Millisecond * 15, + } + pool, err := thriftbp.NewCustomClientPool( + cfg, + thriftbp.SingleAddressGenerator(ln.Addr().String()), + thrift.NewTBinaryProtocolFactoryConf(cfg.ToTConfiguration()), + ) + if err != nil { + t.Fatal(err) + } + + go func() { ln.Accept() }() + pool.TClient().Call(context.Background(), "test", nil, nil) + // conn, _ := ln.Accept() + // buff := make([]byte, 1024) + // conn.Read(buff) + fmt.Println("string(buff)") +} + func TestInitialConnectionsFallback(t *testing.T) { ln, err := net.Listen("tcp", addr) if err != nil { From 091d9c71fa73d1462e8e39e455796749d753b3ed Mon Sep 17 00:00:00 2001 From: Dorian Jaminais-Grellier Date: Mon, 8 Jan 2024 09:36:00 -0800 Subject: [PATCH 6/6] add unit test --- thriftbp/client_pool_test.go | 59 ++++++++++++++++++++++-------------- 1 file changed, 36 insertions(+), 23 deletions(-) diff --git a/thriftbp/client_pool_test.go b/thriftbp/client_pool_test.go index 0650b95d1..f31b09e6e 100644 --- a/thriftbp/client_pool_test.go +++ b/thriftbp/client_pool_test.go @@ -11,8 +11,11 @@ import ( "github.com/apache/thrift/lib/go/thrift" + "github.com/reddit/baseplate.go" "github.com/reddit/baseplate.go/ecinterface" + baseplatethrift "github.com/reddit/baseplate.go/internal/gen-go/reddit/baseplate" "github.com/reddit/baseplate.go/thriftbp" + "github.com/reddit/baseplate.go/thriftbp/thrifttest" ) const ( @@ -190,37 +193,47 @@ func TestBehaviorWithNetworkIssues(t *testing.T) { } } +type thriftHostnameHandler struct { + server baseplate.Server +} + +func (thriftHostnameHandler) IsHealthy(ctx context.Context, _ *baseplatethrift.IsHealthyRequest) (r bool, err error) { + value, ok := thrift.GetHeader(ctx, thriftbp.ThriftHostnameHeader) + if !ok { + return false, errors.New("did not find the thrift header") + } + if value != "my-thrift-header" { + return false, errors.New("unexpected value for the thrift header") + } + return true, nil +} + func TestThriftHostnameHeader(t *testing.T) { - ln, err := net.Listen("tcp", addr) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + store := newSecretsStore(t) + defer store.Close() + + handler := thriftHostnameHandler{} + server, err := thrifttest.NewBaseplateServer(thrifttest.ServerConfig{ + Processor: baseplatethrift.NewBaseplateServiceV2Processor(&handler), + SecretStore: store, + ClientConfig: thriftbp.ClientPoolConfig{ + ThriftHostnameHeader: "my-thrift-header", + }, + }) if err != nil { t.Fatal(err) } - defer ln.Close() + handler.server = server + server.Start(ctx) - cfg := thriftbp.ClientPoolConfig{ - Addr: ":9090", - EdgeContextImpl: ecinterface.Mock(), - ServiceSlug: "test", - InitialConnections: 1, - MaxConnections: 5, - ConnectTimeout: time.Millisecond * 5, - SocketTimeout: time.Millisecond * 15, - } - pool, err := thriftbp.NewCustomClientPool( - cfg, - thriftbp.SingleAddressGenerator(ln.Addr().String()), - thrift.NewTBinaryProtocolFactoryConf(cfg.ToTConfiguration()), - ) + client := baseplatethrift.NewBaseplateServiceV2Client(server.ClientPool.TClient()) + _, err = client.IsHealthy(ctx, &baseplatethrift.IsHealthyRequest{}) if err != nil { t.Fatal(err) } - - go func() { ln.Accept() }() - pool.TClient().Call(context.Background(), "test", nil, nil) - // conn, _ := ln.Accept() - // buff := make([]byte, 1024) - // conn.Read(buff) - fmt.Println("string(buff)") } func TestInitialConnectionsFallback(t *testing.T) {