func TestRelayHandleLargeLocalCall(t *testing.T) {
	opts := testutils.NewOpts().SetRelayOnly().
		SetRelayLocal("relay").
		AddLogFilter("Received fragmented callReq", 1).
		// Expect 4 callReqContinues for 256 kb payload that we cannot relay.
		AddLogFilter("Failed to relay frame.", 4)
	testutils.WithTestServer(t, opts, func(ts *testutils.TestServer) {
		client := ts.NewClient(nil)
		testutils.RegisterEcho(ts.Relay(), nil)

		// This large call should fail with a bad request.
		err := testutils.CallEcho(client, ts.HostPort(), "relay", &raw.Args{
			Arg2: testutils.RandBytes(128 * 1024),
			Arg3: testutils.RandBytes(128 * 1024),
		})
		if assert.Equal(t, ErrCodeBadRequest, GetSystemErrorCode(err), "Expected BadRequest for large call to relay") {
			assert.Contains(t, err.Error(), "cannot receive fragmented calls")
		}

		// We may get an error before the call is finished flushing.
		// Do a ping to ensure everything has been flushed.
		ctx, cancel := NewContext(time.Second)
		defer cancel()
		require.NoError(t, client.Ping(ctx, ts.HostPort()), "Ping failed")
	})
}

func TestLargeRequest(t *testing.T) {
	CheckStress(t)

	const (
		KB = 1024
		MB = 1024 * KB
		GB = 1024 * MB

		maxRequestSize = 1 * GB
	)

	WithVerifiedServer(t, nil, func(serverCh *Channel, hostPort string) {
		serverCh.Register(raw.Wrap(newTestHandler(t)), "echo")

		for reqSize := 2; reqSize <= maxRequestSize; reqSize *= 2 {
			log.Printf("reqSize = %v", reqSize)
			arg3 := testutils.RandBytes(reqSize)
			arg2 := testutils.RandBytes(reqSize / 2)

			clientCh := testutils.NewClient(t, nil)
			ctx, cancel := NewContext(time.Second * 30)
			rArg2, rArg3, _, err := raw.Call(ctx, clientCh, hostPort, serverCh.PeerInfo().ServiceName, "echo", arg2, arg3)
			require.NoError(t, err, "Call failed")

			if !bytes.Equal(arg2, rArg2) {
				t.Errorf("echo arg2 mismatch")
			}
			if !bytes.Equal(arg3, rArg3) {
				t.Errorf("echo arg3 mismatch")
			}
			cancel()
		}
	})
}

func TestDirtyFrameRequests(t *testing.T) {

	argSizes := []int{50000, 100000, 150000}
	WithVerifiedServer(t, &testutils.ChannelOpts{
		ServiceName: "swap-server",
		DefaultConnectionOptions: ConnectionOptions{
			FramePool: dirtyFramePool{},
		},
	}, func(serverCh *Channel, hostPort string) {
		peerInfo := serverCh.PeerInfo()
		serverCh.Register(raw.Wrap(&swapper{t}), "swap")

		for _, arg2Size := range argSizes {
			for _, arg3Size := range argSizes {
				ctx, cancel := NewContext(time.Second)
				defer cancel()

				arg2, arg3 := testutils.RandBytes(arg2Size), testutils.RandBytes(arg3Size)
				res2, res3, _, err := raw.Call(ctx, serverCh, hostPort, peerInfo.ServiceName, "swap", arg2, arg3)
				if assert.NoError(t, err, "Call failed") {
					assert.Equal(t, arg2, res3, "Result arg3 wrong")
					assert.Equal(t, arg3, res2, "Result arg3 wrong")
				}
			}
		}
	})
}

func TestFragmentationSlowReader(t *testing.T) {
	startReading, handlerComplete := make(chan struct{}), make(chan struct{})
	handler := func(ctx context.Context, call *InboundCall) {
		<-ctx.Done()
		<-startReading
		_, err := raw.ReadArgs(call)
		assert.Error(t, err, "ReadArgs should fail since frames will be dropped due to slow reading")
		close(handlerComplete)
	}

	// Inbound forward will timeout and cause a warning log.
	opts := testutils.NewOpts().AddLogFilter("Unable to forward frame", 1)
	WithVerifiedServer(t, opts, func(ch *Channel, hostPort string) {
		ch.Register(HandlerFunc(handler), "echo")

		arg2 := testutils.RandBytes(MaxFramePayloadSize * MexChannelBufferSize)
		arg3 := testutils.RandBytes(MaxFramePayloadSize * (MexChannelBufferSize + 1))

		ctx, cancel := NewContext(testutils.Timeout(15 * time.Millisecond))
		defer cancel()

		_, _, _, err := raw.Call(ctx, ch, hostPort, testServiceName, "echo", arg2, arg3)
		assert.Error(t, err, "Call should timeout due to slow reader")

		close(startReading)
		<-handlerComplete
	})
	goroutines.VerifyNoLeaks(t, nil)
}

func TestFragmentationSlowReader(t *testing.T) {
	startReading, handlerComplete := make(chan struct{}), make(chan struct{})
	handler := func(ctx context.Context, call *InboundCall) {
		<-startReading
		_, err := raw.ReadArgs(call)
		assert.Error(t, err, "ReadArgs should fail since frames will be dropped due to slow reading")
		close(handlerComplete)
	}

	WithVerifiedServer(t, nil, func(ch *Channel, hostPort string) {
		ch.Register(HandlerFunc(handler), "echo")

		arg2 := testutils.RandBytes(MaxFramePayloadSize * MexChannelBufferSize)
		arg3 := testutils.RandBytes(MaxFramePayloadSize * (MexChannelBufferSize + 1))

		ctx, cancel := NewContext(10 * time.Millisecond)
		defer cancel()

		_, _, _, err := raw.Call(ctx, ch, hostPort, testServiceName, "echo", arg2, arg3)
		assert.Error(t, err, "Call should timeout due to slow reader")

		close(startReading)
		<-handlerComplete
	})
	VerifyNoBlockedGoroutines(t)
}

func TestDirtyFrameRequests(t *testing.T) {
	argSizes := []int{25000, 50000, 75000}

	// Create the largest required random cache.
	testutils.RandBytes(argSizes[len(argSizes)-1])

	opts := testutils.NewOpts().
		SetServiceName("swap-server").
		SetFramePool(dirtyFramePool{})
	WithVerifiedServer(t, opts, func(serverCh *Channel, hostPort string) {
		peerInfo := serverCh.PeerInfo()
		serverCh.Register(raw.Wrap(&swapper{t}), "swap")

		for _, argSize := range argSizes {
			ctx, cancel := NewContext(time.Second)
			defer cancel()

			arg2, arg3 := testutils.RandBytes(argSize), testutils.RandBytes(argSize)
			res2, res3, _, err := raw.Call(ctx, serverCh, hostPort, peerInfo.ServiceName, "swap", arg2, arg3)
			if assert.NoError(t, err, "Call failed") {
				assert.Equal(t, arg2, res3, "Result arg3 wrong")
				assert.Equal(t, arg3, res2, "Result arg3 wrong")
			}
		}
	})
}

func TestWriteErrorAfterTimeout(t *testing.T) {
	// TODO: Make this test block at different points (e.g. before, during read/write).
	WithVerifiedServer(t, nil, func(ch *Channel, hostPort string) {
		timedOut := make(chan struct{})
		done := make(chan struct{})
		handler := func(ctx context.Context, call *InboundCall) {
			<-ctx.Done()
			<-timedOut
			_, err := raw.ReadArgs(call)
			assert.Equal(t, ErrTimeout, err, "Read args should fail with timeout")
			response := call.Response()
			assert.Equal(t, ErrTimeout, response.SendSystemError(ErrServerBusy), "SendSystemError should fail")
			close(done)
		}
		ch.Register(HandlerFunc(handler), "call")

		ctx, cancel := NewContext(testutils.Timeout(20 * time.Millisecond))
		defer cancel()
		_, _, _, err := raw.Call(ctx, ch, hostPort, testServiceName, "call", nil, testutils.RandBytes(100000))
		assert.Equal(t, err, ErrTimeout, "Call should timeout")
		close(timedOut)
		<-done
	})
	goroutines.VerifyNoLeaks(t, nil)
}

func getRequestTests(t *testing.T) []requestTest {
	randBytes := testutils.RandBytes(40000)
	return []requestTest{
		{
			name: "get simple",
			f: func(httpAddr string) *http.Request {
				req, err := http.NewRequest("GET", fmt.Sprintf("http://%v/this/is/my?req=1&v=2&v&a&a", httpAddr), nil)
				require.NoError(t, err, "NewRequest failed")
				return req
			},
		},
		{
			name: "post simple",
			f: func(httpAddr string) *http.Request {
				body := strings.NewReader("This is a simple POST body")
				req, err := http.NewRequest("POST", fmt.Sprintf("http://%v/post/path?v=1&b=3", httpAddr), body)
				require.NoError(t, err, "NewRequest failed")
				return req
			},
		},
		{
			name: "post random bytes",
			f: func(httpAddr string) *http.Request {
				body := bytes.NewReader(randBytes)
				req, err := http.NewRequest("POST", fmt.Sprintf("http://%v/post/path?v=1&b=3", httpAddr), body)
				require.NoError(t, err, "NewRequest failed")
				return req
			},
		},
	}
}

func TestRelayMakeOutgoingCall(t *testing.T) {
	opts := testutils.NewOpts().SetRelayOnly()
	testutils.WithTestServer(t, opts, func(ts *testutils.TestServer) {
		svr1 := ts.Relay()
		svr2 := ts.NewServer(testutils.NewOpts().SetServiceName("svc2"))
		testutils.RegisterEcho(svr2, nil)

		sizes := []int{128, 1024, 128 * 1024}
		for _, size := range sizes {
			err := testutils.CallEcho(svr1, ts.HostPort(), "svc2", &raw.Args{
				Arg2: testutils.RandBytes(size),
				Arg3: testutils.RandBytes(size),
			})
			assert.NoError(t, err, "Echo with size %v failed", size)
		}
	})
}

func TestLargeMethod(t *testing.T) {
	WithVerifiedServer(t, nil, func(ch *Channel, hostPort string) {
		ctx, cancel := NewContext(time.Second)
		defer cancel()

		largeMethod := testutils.RandBytes(16*1024 + 1)
		_, _, _, err := raw.Call(ctx, ch, hostPort, testServiceName, string(largeMethod), nil, nil)
		assert.Equal(t, ErrMethodTooLarge, err)
	})
}

func TestLargeMethod(t *testing.T) {
	testutils.WithTestServer(t, nil, func(ts *testutils.TestServer) {
		ctx, cancel := NewContext(time.Second)
		defer cancel()

		largeMethod := testutils.RandBytes(16*1024 + 1)
		_, _, _, err := raw.Call(ctx, ts.Server(), ts.HostPort(), ts.ServiceName(), string(largeMethod), nil, nil)
		assert.Equal(t, ErrMethodTooLarge, err)
	})
}

func TestFragmentationSlowReader(t *testing.T) {
	// Inbound forward will timeout and cause a warning log.
	opts := testutils.NewOpts().
		AddLogFilter("Unable to forward frame", 1).
		AddLogFilter("Connection error", 1)

	testutils.WithTestServer(t, opts, func(ts *testutils.TestServer) {
		startReading, handlerComplete := make(chan struct{}), make(chan struct{})
		handler := func(ctx context.Context, call *InboundCall) {
			<-startReading
			<-ctx.Done()
			_, err := raw.ReadArgs(call)
			assert.Error(t, err, "ReadArgs should fail since frames will be dropped due to slow reading")
			close(handlerComplete)
		}

		ts.Register(HandlerFunc(handler), "echo")

		arg2 := testutils.RandBytes(MaxFramePayloadSize * MexChannelBufferSize)
		arg3 := testutils.RandBytes(MaxFramePayloadSize * (MexChannelBufferSize + 1))

		ctx, cancel := NewContext(testutils.Timeout(30 * time.Millisecond))
		defer cancel()

		_, _, _, err := raw.Call(ctx, ts.Server(), ts.HostPort(), ts.ServiceName(), "echo", arg2, arg3)
		assert.Error(t, err, "Call should timeout due to slow reader")

		close(startReading)
		select {
		case <-handlerComplete:
		case <-time.After(testutils.Timeout(70 * time.Millisecond)):
			t.Errorf("Handler not called, context timeout may be too low")
		}

		calls := relaytest.NewMockStats()
		calls.Add(ts.ServiceName(), ts.ServiceName(), "echo").Failed("timeout").End()
		ts.AssertRelayStats(calls)
	})
}

func doPingAndCall(t *testing.T, clientCh *Channel, hostPort string) {
	ctx, cancel := NewContext(time.Second * 5)
	defer cancel()

	require.NoError(t, clientCh.Ping(ctx, hostPort))

	const maxRandArg = 512 * 1024

	arg2 := testutils.RandBytes(rand.Intn(maxRandArg))
	arg3 := testutils.RandBytes(rand.Intn(maxRandArg))
	resArg2, resArg3, _, err := raw.Call(ctx, clientCh, hostPort, "swap-server", "swap", arg2, arg3)
	if !assert.NoError(t, err, "error during sendRecv") {
		return
	}

	// We expect the arguments to be swapped.
	if bytes.Compare(arg3, resArg2) != 0 {
		t.Errorf("returned arg2 does not match expected:\n  got %v\n want %v", resArg2, arg3)
	}
	if bytes.Compare(arg2, resArg3) != 0 {
		t.Errorf("returned arg2 does not match expected:\n  got %v\n want %v", resArg3, arg2)
	}
}

func TestWriteArg3AfterTimeout(t *testing.T) {
	// The channel reads and writes during timeouts, causing warning logs.
	opts := testutils.NewOpts().DisableLogVerification()
	testutils.WithTestServer(t, opts, func(ts *testutils.TestServer) {
		timedOut := make(chan struct{})

		handler := func(ctx context.Context, call *InboundCall) {
			_, err := raw.ReadArgs(call)
			assert.NoError(t, err, "Read args failed")
			response := call.Response()
			assert.NoError(t, NewArgWriter(response.Arg2Writer()).Write(nil), "Write Arg2 failed")
			writer, err := response.Arg3Writer()
			assert.NoError(t, err, "Arg3Writer failed")

			for {
				if _, err := writer.Write(testutils.RandBytes(4096)); err != nil {
					assert.Equal(t, err, ErrTimeout, "Handler should timeout")
					close(timedOut)
					return
				}
				runtime.Gosched()
			}
		}
		ts.Register(HandlerFunc(handler), "call")

		ctx, cancel := NewContext(testutils.Timeout(50 * time.Millisecond))
		defer cancel()

		_, _, _, err := raw.Call(ctx, ts.Server(), ts.HostPort(), ts.ServiceName(), "call", nil, nil)
		assert.Equal(t, err, ErrTimeout, "Call should timeout")

		// Wait for the write to complete, make sure there's no errors.
		select {
		case <-time.After(testutils.Timeout(30 * time.Millisecond)):
			t.Errorf("Handler should have failed due to timeout")
		case <-timedOut:
		}

		calls := relaytest.NewMockStats()
		calls.Add(ts.ServiceName(), ts.ServiceName(), "call").Failed("timeout").Succeeded().End()
		ts.AssertRelayStats(calls)
	})
}

func TestWriteArg3AfterTimeout(t *testing.T) {
	// The channel reads and writes during timeouts, causing warning logs.
	opts := testutils.NewOpts().DisableLogVerification()
	WithVerifiedServer(t, opts, func(ch *Channel, hostPort string) {
		timedOut := make(chan struct{})

		handler := func(ctx context.Context, call *InboundCall) {
			_, err := raw.ReadArgs(call)
			assert.NoError(t, err, "Read args failed")
			response := call.Response()
			assert.NoError(t, NewArgWriter(response.Arg2Writer()).Write(nil), "Write Arg2 failed")
			writer, err := response.Arg3Writer()
			assert.NoError(t, err, "Arg3Writer failed")

			for {
				if _, err := writer.Write(testutils.RandBytes(4096)); err != nil {
					assert.Equal(t, err, ErrTimeout, "Handler should timeout")
					close(timedOut)
					return
				}
				runtime.Gosched()
			}
		}
		ch.Register(HandlerFunc(handler), "call")

		ctx, cancel := NewContext(20 * time.Millisecond)
		defer cancel()
		_, _, _, err := raw.Call(ctx, ch, hostPort, testServiceName, "call", nil, nil)
		assert.Equal(t, err, ErrTimeout, "Call should timeout")

		// Wait for the write to complete, make sure there's no errors.
		select {
		case <-time.After(30 * time.Millisecond):
			t.Errorf("Handler should have failed due to timeout")
		case <-timedOut:
		}
	})
	VerifyNoBlockedGoroutines(t)
}

func TestFramesReleased(t *testing.T) {
	CheckStress(t)

	defer testutils.SetTimeout(t, 10*time.Second)()
	const (
		requestsPerGoroutine = 10
		numGoroutines        = 10
		maxRandArg           = 512 * 1024
	)

	var serverExchanges, clientExchanges string
	pool := NewRecordingFramePool()
	opts := testutils.NewOpts().
		SetServiceName("swap-server").
		SetFramePool(pool)
	WithVerifiedServer(t, opts, func(serverCh *Channel, hostPort string) {
		serverCh.Register(raw.Wrap(&swapper{t}), "swap")

		clientCh, err := NewChannel("swap-client", nil)
		require.NoError(t, err)
		defer clientCh.Close()

		// Create an active connection that can be shared by the goroutines by calling Ping.
		ctx, cancel := NewContext(time.Second)
		defer cancel()
		require.NoError(t, clientCh.Ping(ctx, hostPort))

		var wg sync.WaitGroup
		worker := func() {
			for i := 0; i < requestsPerGoroutine; i++ {
				ctx, cancel := NewContext(time.Second * 5)
				defer cancel()

				require.NoError(t, clientCh.Ping(ctx, hostPort))

				arg2 := testutils.RandBytes(rand.Intn(maxRandArg))
				arg3 := testutils.RandBytes(rand.Intn(maxRandArg))
				resArg2, resArg3, _, err := raw.Call(ctx, clientCh, hostPort, "swap-server", "swap", arg2, arg3)
				if !assert.NoError(t, err, "error during sendRecv") {
					continue
				}

				// We expect the arguments to be swapped.
				if bytes.Compare(arg3, resArg2) != 0 {
					t.Errorf("returned arg2 does not match expected:\n  got %v\n want %v", resArg2, arg3)
				}
				if bytes.Compare(arg2, resArg3) != 0 {
					t.Errorf("returned arg2 does not match expected:\n  got %v\n want %v", resArg3, arg2)
				}
			}
			wg.Done()
		}

		for i := 0; i < numGoroutines; i++ {
			wg.Add(1)
			go worker()
		}

		wg.Wait()

		serverExchanges = CheckEmptyExchanges(serverCh)
		clientExchanges = CheckEmptyExchanges(clientCh)
	})

	// Wait a few milliseconds for the closing of channels to take effect.
	time.Sleep(10 * time.Millisecond)

	if unreleasedCount, isEmpty := pool.CheckEmpty(); isEmpty != "" || unreleasedCount > 0 {
		t.Errorf("Frame pool has %v unreleased frames, errors:\n%v", unreleasedCount, isEmpty)
	}

	// Check the message exchanges and make sure they are all empty.
	if serverExchanges != "" {
		t.Errorf("Found uncleared message exchanges on server:\n%s", serverExchanges)
	}
	if clientExchanges != "" {
		t.Errorf("Found uncleared message exchanges on client:\n%s", clientExchanges)
	}
}

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 TestWriteErrorAfterTimeout(t *testing.T) {
	// TODO: Make this test block at different points (e.g. before, during read/write).
	testutils.WithTestServer(t, nil, func(ts *testutils.TestServer) {
		timedOut := make(chan struct{})
		done := make(chan struct{})
		handler := func(ctx context.Context, call *InboundCall) {
			<-ctx.Done()
			<-timedOut
			_, err := raw.ReadArgs(call)
			assert.Equal(t, ErrTimeout, err, "Read args should fail with timeout")
			response := call.Response()
			assert.Equal(t, ErrTimeout, response.SendSystemError(ErrServerBusy), "SendSystemError should fail")
			close(done)
		}
		ts.Register(HandlerFunc(handler), "call")

		ctx, cancel := NewContext(testutils.Timeout(30 * time.Millisecond))
		defer cancel()
		_, _, _, err := raw.Call(ctx, ts.Server(), ts.HostPort(), ts.ServiceName(), "call", nil, testutils.RandBytes(100000))
		assert.Equal(t, err, ErrTimeout, "Call should timeout")
		close(timedOut)

		select {
		case <-done:
		case <-time.After(time.Second):
			t.Errorf("Handler not called, timeout may be too low")
		}

		calls := relaytest.NewMockStats()
		calls.Add(ts.ServiceName(), ts.ServiceName(), "call").Failed("timeout").End()
		ts.AssertRelayStats(calls)
	})
}