Skip to content

Commit

Permalink
always update thrift header write list in headerbp client middleware (#…
Browse files Browse the repository at this point in the history
…685)

This is really just to make the code cleaner since it technically works because thrift.GetWriteHeaderList returns the underlying list, not a copy, and slices.Delete deletes from the source slice. I think the code is a lot clearer and less likely to have a bug later if we always set the outgoing header list.

I verified this by adding a test case. The test passes without my changes, but it fails if I copy outgoing before deleting the headers from it.
  • Loading branch information
pacejackson authored Feb 20, 2025
1 parent cc43adf commit ddd6634
Show file tree
Hide file tree
Showing 2 changed files with 54 additions and 1 deletion.
2 changes: 1 addition & 1 deletion thriftbp/client_middlewares.go
Original file line number Diff line number Diff line change
Expand Up @@ -449,8 +449,8 @@ func ClientBaseplateHeadersMiddleware(service, client string) thrift.ClientMiddl
for k, v := range toAdd {
ctx = thrift.SetHeader(ctx, k, v)
}
ctx = thrift.SetWriteHeaderList(ctx, outgoing)
}
ctx = thrift.SetWriteHeaderList(ctx, outgoing)
return next.Call(ctx, method, args, result)
},
}
Expand Down
53 changes: 53 additions & 0 deletions thriftbp/server_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -141,6 +141,59 @@ func TestHeaderPropagation(t *testing.T) {
}
}

// verify that new headers are removed when there are no headers set to propagate
func TestHeaderPropagation_removeOnly(t *testing.T) {
store := newSecretsStore(t)

ctx, cancel := context.WithCancel(context.Background())
t.Cleanup(cancel)

ecImpl := ecinterface.Mock()

downstreamProcessor := baseplatethrift.NewBaseplateServiceV2Processor(&headerPropagationVerificationService{
wantUnset: []string{"x-bp-test"},
})
downstreamServer, err := thrifttest.NewBaseplateServer(thrifttest.ServerConfig{
Processor: downstreamProcessor,
SecretStore: store,
EdgeContextImpl: ecImpl,
})
if err != nil {
t.Fatal(err)
}
downstreamServer.Start(ctx)
time.Sleep(100 * time.Millisecond) // wait for the server to start

originProcessor := baseplatethrift.NewBaseplateServiceV2Processor(&headerPropagationVerificationService{
wantUnset: []string{"x-bp-test"},
client: func() baseplatethrift.BaseplateServiceV2 {
return baseplatethrift.NewBaseplateServiceV2Client(downstreamServer.ClientPool.TClient())
},
})
server, err := thrifttest.NewBaseplateServer(thrifttest.ServerConfig{
Processor: originProcessor,
SecretStore: store,
})
if err != nil {
t.Fatal(err)
}
server.Start(ctx)
time.Sleep(100 * time.Millisecond) // wait for the server to start

client := baseplatethrift.NewBaseplateServiceV2Client(server.ClientPool.TClient())
got, err := client.IsHealthy(ctx, &baseplatethrift.IsHealthyRequest{
Probe: baseplatethrift.IsHealthyProbePtr(baseplatethrift.IsHealthyProbe_READINESS),
})

if err != nil {
t.Errorf("expected no error, got %v", err)
}

const want = true
if got != want {
t.Errorf("success mismatch, want %v, got %v", want, got)
}
}
func setHeader(ctx context.Context, key, value string) context.Context {
ctx = thrift.SetHeader(ctx, key, value)
return thrift.SetWriteHeaderList(ctx, append(thrift.GetWriteHeaderList(ctx), key))
Expand Down

0 comments on commit ddd6634

Please sign in to comment.