func TestRelay(t *testing.T) {
withRelayedEcho(t, func(_, _, client *Channel, ts *testutils.TestServer) {
tests := []struct {
header string
body string
}{
{"fake-header", "fake-body"}, // fits in one frame
{"fake-header", strings.Repeat("fake-body", 10000)}, // requires continuation
}
sc := client.GetSubChannel("test")
for _, tt := range tests {
ctx, cancel := NewContext(time.Second)
defer cancel()
arg2, arg3, _, err := raw.CallSC(ctx, sc, "echo", []byte(tt.header), []byte(tt.body))
require.NoError(t, err, "Relayed call failed.")
assert.Equal(t, tt.header, string(arg2), "Header was mangled during relay.")
assert.Equal(t, tt.body, string(arg3), "Body was mangled during relay.")
}
calls := relaytest.NewMockStats()
peer := relay.Peer{HostPort: ts.Server().PeerInfo().HostPort}
for _ = range tests {
calls.Add("client", "test", "echo").SetPeer(peer).Succeeded().End()
}
ts.AssertRelayStats(calls)
})
}
func TestMockIgnoresDown(t *testing.T) {
mockHB, err := mockhyperbahn.New()
require.NoError(t, err, "Failed to set up mock hyperbahm")
var (
moe1Called atomic.Bool
moe2Called atomic.Bool
)
moe1 := newAdvertisedEchoServer(t, "moe", mockHB, func() { moe1Called.Store(true) })
defer moe1.Close()
moe2 := newAdvertisedEchoServer(t, "moe", mockHB, func() { moe2Called.Store(true) })
defer moe2.Close()
client := newAdvertisedEchoServer(t, "client", mockHB, nil)
ctx, cancel := tchannel.NewContext(time.Second)
defer cancel()
for i := 0; i < 20; i++ {
_, _, _, err = raw.CallSC(ctx, client.GetSubChannel("moe"), "echo", nil, nil)
assert.NoError(t, err, "Call failed")
}
require.True(t, moe1Called.Load(), "moe1 not called")
require.True(t, moe2Called.Load(), "moe2 not called")
// If moe2 is brought down, all calls should now be sent to moe1.
moe2.Close()
// Wait for the mock HB to have 0 connections to moe
ok := testutils.WaitFor(time.Second, func() bool {
in, out := mockHB.Channel().Peers().GetOrAdd(moe2.PeerInfo().HostPort).NumConnections()
return in+out == 0
})
require.True(t, ok, "Failed waiting for mock HB to have 0 connections")
// Make sure that all calls succeed (they should all go to moe2)
moe1Called.Store(false)
moe2Called.Store(false)
for i := 0; i < 20; i++ {
_, _, _, err = raw.CallSC(ctx, client.GetSubChannel("moe"), "echo", nil, nil)
assert.NoError(t, err, "Call failed")
}
require.True(t, moe1Called.Load(), "moe1 not called")
require.False(t, moe2Called.Load(), "moe2 should not be called after Close")
}
func TestRelayHandlesClosedPeers(t *testing.T) {
opts := serviceNameOpts("test").SetRelayOnly().
// Disable logs as we are closing connections that can error in a lot of places.
DisableLogVerification()
testutils.WithTestServer(t, opts, func(ts *testutils.TestServer) {
ctx, cancel := NewContext(300 * time.Millisecond)
defer cancel()
testutils.RegisterEcho(ts.Server(), nil)
client := ts.NewClient(serviceNameOpts("client"))
client.Peers().Add(ts.HostPort())
sc := client.GetSubChannel("test")
_, _, _, err := raw.CallSC(ctx, sc, "echo", []byte("fake-header"), []byte("fake-body"))
require.NoError(t, err, "Relayed call failed.")
ts.Server().Close()
require.NotPanics(t, func() {
raw.CallSC(ctx, sc, "echo", []byte("fake-header"), []byte("fake-body"))
})
})
}
func TestMockForwards(t *testing.T) {
mockHB, err := mockhyperbahn.New()
require.NoError(t, err, "Failed to set up mock hyperbahm")
called := false
server := newAdvertisedEchoServer(t, "svr", mockHB, func() {
called = true
})
defer server.Close()
client := newAdvertisedEchoServer(t, "client", mockHB, nil)
defer client.Close()
ctx, cancel := tchannel.NewContext(time.Second)
defer cancel()
_, _, _, err = raw.CallSC(ctx, client.GetSubChannel("svr"), "echo", nil, nil)
require.NoError(t, err, "Call failed")
require.True(t, called, "Advertised server was not called")
}
func (c *internalClient) RawCallBuffer(latencies []time.Duration) error {
return c.makeCalls(latencies, func() (time.Duration, error) {
ctx, cancel := tchannel.NewContext(c.opts.timeout)
defer cancel()
started := time.Now()
rArg2, rArg3, _, err := raw.CallSC(ctx, c.sc, "echo", c.argBytes, c.argBytes)
duration := time.Since(started)
if err != nil {
return 0, err
}
if c.checkResult {
if !bytes.Equal(rArg2, c.argBytes) || !bytes.Equal(rArg3, c.argBytes) {
fmt.Println("Arg2", rArg2, "Expect", c.argBytes)
fmt.Println("Arg3", rArg3, "Expect", c.argBytes)
panic("echo call returned wrong results")
}
}
return duration, nil
})
}
func (pt *peerSelectionTest) makeCall(sc *SubChannel) {
ctx, cancel := NewContext(time.Second)
defer cancel()
_, _, _, err := raw.CallSC(ctx, sc, "echo", nil, nil)
assert.NoError(pt.t, err, "raw.Call failed")
}
func benchmarkCallsN(b *testing.B, c benchmarkConfig) {
var (
clients []*Channel
servers []*Channel
)
lt := newLatencyTracker()
if c.numBytes == 0 {
c.numBytes = 100
}
data := testutils.RandBytes(c.numBytes)
// Set up clients and servers.
for i := 0; i < c.numServers; i++ {
servers = append(servers, setupServer(b))
}
for i := 0; i < c.numClients; i++ {
clients = append(clients, testutils.NewClient(b, nil))
for _, s := range servers {
clients[i].Peers().Add(s.PeerInfo().HostPort)
// Initialize a connection
ctx, cancel := NewContext(50 * time.Millisecond)
assert.NoError(b, clients[i].Ping(ctx, s.PeerInfo().HostPort), "Initial ping failed")
cancel()
}
}
// Make calls from clients to the servers
call := func(sc *SubChannel) {
ctx, cancel := NewContext(50 * time.Millisecond)
start := time.Now()
_, _, _, err := raw.CallSC(ctx, sc, "echo", nil, data)
duration := time.Since(start)
cancel()
if assert.NoError(b, err, "Call failed") {
lt.addLatency(duration)
}
}
reqsLeft := testutils.Decrementor(c.numCalls)
clientWorker := func(client *Channel, clientNum, workerNum int) {
sc := client.GetSubChannel(benchService)
for reqsLeft() {
call(sc)
}
}
clientRunner := func(client *Channel, clientNum int) {
testutils.RunN(c.workersPerClient, func(i int) {
clientWorker(client, clientNum, i)
})
}
lt.reset()
defer lt.report(b)
b.ResetTimer()
testutils.RunN(c.numClients, func(i int) {
clientRunner(clients[i], i)
})
}
func TestSetHandler(t *testing.T) {
// Generate a Handler that expects only the given methods to be called.
genHandler := func(methods ...string) Handler {
allowedMethods := make(map[string]struct{}, len(methods))
for _, m := range methods {
allowedMethods[m] = struct{}{}
}
return HandlerFunc(func(ctx context.Context, call *InboundCall) {
method := call.MethodString()
assert.Contains(t, allowedMethods, method, "unexpected call to %q", method)
err := raw.WriteResponse(call.Response(), &raw.Res{Arg3: []byte(method)})
require.NoError(t, err)
})
}
ch := testutils.NewServer(t, testutils.NewOpts().
AddLogFilter("Couldn't find handler", 1, "serviceName", "svc2", "method", "bar"))
defer ch.Close()
// Catch-all handler for the main channel that accepts foo, bar, and baz,
// and a single registered handler for a different subchannel.
ch.GetSubChannel("svc1").SetHandler(genHandler("foo", "bar", "baz"))
ch.GetSubChannel("svc2").Register(genHandler("foo"), "foo")
client := testutils.NewClient(t, nil)
client.Peers().Add(ch.PeerInfo().HostPort)
defer client.Close()
tests := []struct {
Service string
Method string
ShouldFail bool
}{
{"svc1", "foo", false},
{"svc1", "bar", false},
{"svc1", "baz", false},
{"svc2", "foo", false},
{"svc2", "bar", true},
}
for _, tt := range tests {
c := client.GetSubChannel(tt.Service)
ctx, _ := NewContext(time.Second)
_, data, _, err := raw.CallSC(ctx, c, tt.Method, nil, []byte("irrelevant"))
if tt.ShouldFail {
require.Error(t, err)
} else {
require.NoError(t, err)
assert.Equal(t, tt.Method, string(data))
}
}
st := ch.IntrospectState(nil)
assert.Equal(t, "overriden", st.SubChannels["svc1"].Handler.Type.String())
assert.Nil(t, st.SubChannels["svc1"].Handler.Methods)
assert.Equal(t, "methods", st.SubChannels["svc2"].Handler.Type.String())
assert.Equal(t, []string{"foo"}, st.SubChannels["svc2"].Handler.Methods)
}