// 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.

// brpc - A framework to host and access services throughout Baidu.

// Date: 2015/10/22 16:28:44

#include <gtest/gtest.h>
#include <atomic>
#include "brpc/server.h"

#include "brpc/controller.h"
#include "brpc/channel.h"
#include "brpc/callback.h"
#include "brpc/socket.h"
#include "brpc/stream_impl.h"
#include "brpc/policy/streaming_rpc_protocol.h"
#include "echo.pb.h"

class AfterAcceptStream {
public:
    virtual void action(brpc::StreamId) = 0;
};

class MyServiceWithStream : public test::EchoService {
public:
    MyServiceWithStream(const brpc::StreamOptions& options) 
        : _options(options)
        , _after_accept_stream(NULL)
    {}
    MyServiceWithStream(const brpc::StreamOptions& options,
                        AfterAcceptStream* after_accept_stream) 
        : _options(options)
        , _after_accept_stream(after_accept_stream)
    {}
    MyServiceWithStream()
        : _options()
        , _after_accept_stream(NULL)
    {}

    void Echo(::google::protobuf::RpcController* controller,
                       const ::test::EchoRequest* request,
                       ::test::EchoResponse* response,
                       ::google::protobuf::Closure* done) {
        brpc::ClosureGuard done_guard(done);
        response->set_message(request->message());
        brpc::Controller* cntl = (brpc::Controller*)controller;
        brpc::StreamId response_stream;
        ASSERT_EQ(0, StreamAccept(&response_stream, *cntl, &_options));
        LOG(INFO) << "Created response_stream=" << response_stream;
        if (_after_accept_stream) {
            _after_accept_stream->action(response_stream);
        }
    }
private:
    brpc::StreamOptions _options;
    AfterAcceptStream* _after_accept_stream;
};

class StreamingRpcTest : public testing::Test {
protected:
    void SetUp() { request.set_message("hello world"); }
    void TearDown() {}

    test::EchoRequest request;
    test::EchoResponse response;
};

struct BatchStreamFeedbackRaceState {
    brpc::StreamId server_first_stream_id{brpc::INVALID_STREAM_ID};
    brpc::StreamId server_extra_stream_id{brpc::INVALID_STREAM_ID};
    brpc::StreamId client_extra_stream_id{brpc::INVALID_STREAM_ID};

    std::atomic<int> server_first_write_rc{-1};
    std::atomic<int> server_second_write_rc{-1};
    std::atomic<bool> client_got_first_msg{false};
    std::atomic<bool> client_got_second_msg{false};
    std::atomic<bool> server_write_done{false};
    std::atomic<bool> rpc_done{false};
    std::atomic<int> client_closed_count{0};

    bthread_t server_send_tid{0};
    std::atomic<bool> server_send_started{false};
};

class BatchStreamClientHandler : public brpc::StreamInputHandler {
public:
    explicit BatchStreamClientHandler(BatchStreamFeedbackRaceState* state)
        : _state(state) {}

    int on_received_messages(brpc::StreamId id,
                             butil::IOBuf* const messages[],
                             size_t size) override {
        if (id != _state->client_extra_stream_id) {
            // This test only cares about extra stream in batch creation.
            return 0;
        }
        for (size_t i = 0; i < size; ++i) {
            const size_t len = messages[i]->length();
            messages[i]->clear();
            // First payload: 64 bytes. Second payload: 1 byte.
            if (len == 64) {
                _state->client_got_first_msg.store(true, std::memory_order_release);
            } else if (len == 1) {
                _state->client_got_second_msg.store(true, std::memory_order_release);
            }
        }
        return 0;
    }

    void on_idle_timeout(brpc::StreamId /*id*/) override {}

    void on_closed(brpc::StreamId /*id*/) override {
        _state->client_closed_count.fetch_add(1, std::memory_order_release);
    }

    void on_failed(brpc::StreamId /*id*/, int /*error_code*/, const std::string& /*error_text*/) override {}

private:
    BatchStreamFeedbackRaceState* _state;
};

static void* SendTwoMessagesOnServerExtraStream(void* arg) {
    auto* state = static_cast<BatchStreamFeedbackRaceState*>(arg);
    const brpc::StreamId sid = state->server_extra_stream_id;

    // Wait until server-side stream is connected.
    const int64_t connect_deadline_us = butil::gettimeofday_us() + 2 * 1000 * 1000L;
    bool connected = false;
    while (butil::gettimeofday_us() < connect_deadline_us) {
        brpc::SocketUniquePtr ptr;
        if (brpc::Socket::Address(sid, &ptr) == 0) {
            brpc::Stream* s = static_cast<brpc::Stream*>(ptr->conn());
            if (s->_host_socket != NULL && s->_connected) {
                connected = true;
                break;
            }
        }
        usleep(1000);
    }

    if (!connected) {
        state->server_first_write_rc.store(ETIMEDOUT, std::memory_order_relaxed);
        state->server_second_write_rc.store(ETIMEDOUT, std::memory_order_relaxed);
        state->server_write_done.store(true, std::memory_order_release);
        return NULL;
    }

    // 1) Send a payload exactly equal to max_buf_size(64).
    {
        std::string payload(64, 'a');
        butil::IOBuf out;
        out.append(payload);
        state->server_first_write_rc.store(brpc::StreamWrite(sid, out), std::memory_order_relaxed);
    }

    // 2) Then send another byte. This write should become writable only after
    // client sends FEEDBACK with consumed_size >= 64.
    const int64_t write_deadline_us = butil::gettimeofday_us() + 2 * 1000 * 1000L;
    int rc = -1;
    while (butil::gettimeofday_us() < write_deadline_us) {
        butil::IOBuf out;
        out.append("b", 1);
        rc = brpc::StreamWrite(sid, out);
        if (rc == 0) {
            break;
        }
        if (rc != EAGAIN) {
            break;
        }
        const timespec duetime = butil::milliseconds_from_now(100);
        (void)brpc::StreamWait(sid, &duetime);
    }
    state->server_second_write_rc.store(rc, std::memory_order_relaxed);
    state->server_write_done.store(true, std::memory_order_release);
    return NULL;
}

class MyServiceWithBatchStream : public test::EchoService {
public:
    MyServiceWithBatchStream(const brpc::StreamOptions& options,
                             BatchStreamFeedbackRaceState* state)
        : _options(options), _state(state) {}

    void Echo(::google::protobuf::RpcController* controller,
              const ::test::EchoRequest* request,
              ::test::EchoResponse* response,
              ::google::protobuf::Closure* done) override {
        brpc::ClosureGuard done_guard(done);
        response->set_message(request->message());
        brpc::Controller* cntl = static_cast<brpc::Controller*>(controller);

        brpc::StreamIds response_streams;
        ASSERT_EQ(0, brpc::StreamAccept(response_streams, *cntl, &_options));
        ASSERT_EQ(2u, response_streams.size());
        _state->server_first_stream_id = response_streams[0];
        _state->server_extra_stream_id = response_streams[1];

        bthread_t tid;
        ASSERT_EQ(0, bthread_start_background(
                         &tid, &BTHREAD_ATTR_NORMAL,
                         SendTwoMessagesOnServerExtraStream, _state));
        _state->server_send_tid = tid;
        _state->server_send_started.store(true, std::memory_order_release);
    }

private:
    brpc::StreamOptions _options;
    BatchStreamFeedbackRaceState* _state;
};

static void SetAtomicTrue(std::atomic<bool>* f) {
    f->store(true, std::memory_order_release);
}

template <typename Pred>
static bool WaitForTrue(Pred pred, int timeout_ms) {
    const int64_t deadline_us = butil::gettimeofday_us() + (int64_t)timeout_ms * 1000L;
    while (!pred() && butil::gettimeofday_us() < deadline_us) {
        usleep(1000);
    }
    return pred();
}

static bool WaitForTrue(const std::atomic<bool>& f, int timeout_ms) {
    return WaitForTrue([&f]() { return f.load(std::memory_order_acquire); }, timeout_ms);
}

TEST_F(StreamingRpcTest, sanity) {
    brpc::Server server;
    MyServiceWithStream service;
    ASSERT_EQ(0, server.AddService(&service, brpc::SERVER_DOESNT_OWN_SERVICE));
    ASSERT_EQ(0, server.Start(9007, NULL));
    brpc::Channel channel;
    ASSERT_EQ(0, channel.Init("127.0.0.1:9007", NULL));
    brpc::Controller cntl;
    brpc::StreamId request_stream;
    ASSERT_EQ(0, StreamCreate(&request_stream, cntl, NULL));
    brpc::ScopedStream stream_guard(request_stream);
    test::EchoService_Stub stub(&channel);
    stub.Echo(&cntl, &request, &response, NULL);
    ASSERT_FALSE(cntl.Failed()) << cntl.ErrorText() << " request_stream=" << request_stream;
    usleep(10);
    brpc::StreamClose(request_stream);
    server.Stop(0);
    server.Join();
}

TEST_F(StreamingRpcTest, batch_create_stream_feedback_race) {
    BatchStreamFeedbackRaceState state;
    BatchStreamClientHandler client_handler(&state);

    brpc::StreamOptions server_stream_opt;
    // Make server-side sender sensitive to FEEDBACK quickly.
    server_stream_opt.max_buf_size = 16;

    brpc::Server server;
    MyServiceWithBatchStream service(server_stream_opt, &state);
    ASSERT_EQ(0, server.AddService(&service, brpc::SERVER_DOESNT_OWN_SERVICE));
    ASSERT_EQ(0, server.Start(9007, NULL));

    brpc::Channel channel;
    ASSERT_EQ(0, channel.Init("127.0.0.1:9007", NULL));

    brpc::Controller cntl;
    brpc::StreamIds request_streams;
    brpc::StreamOptions client_stream_opt;
    client_stream_opt.handler = &client_handler;
    client_stream_opt.max_buf_size = 0;
    ASSERT_EQ(0, brpc::StreamCreate(request_streams, 2, cntl, &client_stream_opt));
    ASSERT_EQ(2u, request_streams.size());
    state.client_extra_stream_id = request_streams[1];

    // Block SetConnected() on the extra stream to enlarge the race window.
    brpc::SocketUniquePtr client_extra_ptr;
    ASSERT_EQ(0, brpc::Socket::Address(state.client_extra_stream_id, &client_extra_ptr));
    brpc::Stream* client_extra_stream = static_cast<brpc::Stream*>(client_extra_ptr->conn());
    bthread_mutex_lock(&client_extra_stream->_connect_mutex);
    struct UnlockGuard {
        bthread_mutex_t* m;
        ~UnlockGuard() {
            if (m) {
                bthread_mutex_unlock(m);
            }
        }
    } unlock_guard{&client_extra_stream->_connect_mutex};

    BRPC_SCOPE_EXIT {
        if (state.server_extra_stream_id != brpc::INVALID_STREAM_ID) {
            brpc::StreamClose(state.server_extra_stream_id);
        }
        if (state.server_first_stream_id != brpc::INVALID_STREAM_ID) {
            brpc::StreamClose(state.server_first_stream_id);
        }
        for (auto sid : request_streams) {
            brpc::StreamClose(sid);
        }

        if (state.server_send_tid) {
            bthread_join(state.server_send_tid, NULL);
        }
        server.Stop(0);
        server.Join();

        // Release the SocketUniquePtr held above so the fake socket can be
        // recycled. Otherwise BeforeRecycle / on_closed for the extra stream
        // is deferred until `client_extra_ptr` destructs at scope exit, which
        // happens *after* `client_handler` and `state` are destroyed -> UAF
        // inside Stream::Consume on Linux.
        client_extra_ptr.reset();

        // on_closed() runs asynchronously on each client stream's consumer
        // bthread. Wait for both before letting handler/state go out of
        // scope, otherwise Stream::Consume will dereference freed memory.
        int expected_closed = request_streams.size();
        WaitForTrue([&state, expected_closed]() {
            return state.client_closed_count.load(std::memory_order_acquire)
                   >= expected_closed;
        }, 2000);
    };

    test::EchoService_Stub stub(&channel);
    stub.Echo(&cntl, &request, &response, brpc::NewCallback(SetAtomicTrue, &state.rpc_done));

    // Wait until client consumes the first 64B payload on extra stream.
    ASSERT_TRUE(WaitForTrue(state.client_got_first_msg, 2000));

    // Unblock SetConnected(); the fix in PR 3215 should send the first FEEDBACK
    // with consumed_size=64 here, making server-side stream writable again.
    bthread_mutex_unlock(&client_extra_stream->_connect_mutex);
    unlock_guard.m = NULL;

    ASSERT_TRUE(WaitForTrue(state.rpc_done, 2000));
    ASSERT_FALSE(cntl.Failed()) << cntl.ErrorText();

    // Wait for server-side send thread to be started.
    ASSERT_TRUE(WaitForTrue(state.server_send_started, 2000));

    ASSERT_TRUE(WaitForTrue(state.server_write_done, 2000));
    ASSERT_EQ(0, state.server_first_write_rc.load(std::memory_order_relaxed));
    ASSERT_EQ(0, state.server_second_write_rc.load(std::memory_order_relaxed));
    ASSERT_TRUE(WaitForTrue(state.client_got_second_msg, 2000));
}

struct HandlerControl {
    HandlerControl()
        : block(false)
    {}
    bool block;
};

class OrderedInputHandler : public brpc::StreamInputHandler {
public:
    explicit OrderedInputHandler(HandlerControl *cntl = NULL)
        : _expected_next_value(0)
        , _failed(false)
        , _stopped(false)
        , _idle_times(0)
        , _cntl(cntl)
    {}

    int on_received_messages(brpc::StreamId /*id*/,
                             butil::IOBuf *const messages[],
                             size_t size) override {
        if (_cntl && _cntl->block) {
            while (_cntl->block) {
                usleep(100);
            }
        }
        for (size_t i = 0; i < size; ++i) {
            CHECK(messages[i]->length() == sizeof(int));
            int network = 0;
            messages[i]->cutn(&network, sizeof(int));
            EXPECT_EQ((int)ntohl(network), _expected_next_value++);
        }
        return 0;
    }

    void on_idle_timeout(brpc::StreamId /*id*/) override {
        ++_idle_times;
    }

    void on_closed(brpc::StreamId /*id*/) override {
        ASSERT_FALSE(_stopped);
        _stopped = true;
    }

    void on_failed(brpc::StreamId id, int error_code,
                   const std::string& /*error_text*/) override {
        ASSERT_FALSE(_failed);
        ASSERT_NE(0, error_code);
        _failed = true;
    }

    bool failed() const { return _failed; }
    bool stopped() const { return _stopped; }
    int idle_times() const { return _idle_times; }
private:
    int _expected_next_value;
    bool _failed;
    bool _stopped;
    int _idle_times;
    HandlerControl* _cntl;
};

TEST_F(StreamingRpcTest, received_in_order) {
    OrderedInputHandler handler;
    brpc::StreamOptions opt;
    opt.handler = &handler;
    opt.messages_in_batch = 100;
    brpc::Server server;
    MyServiceWithStream service(opt);
    ASSERT_EQ(0, server.AddService(&service, brpc::SERVER_DOESNT_OWN_SERVICE));
    ASSERT_EQ(0, server.Start(9007, NULL));
    brpc::Channel channel;
    ASSERT_EQ(0, channel.Init("127.0.0.1:9007", NULL));
    brpc::Controller cntl;
    brpc::StreamId request_stream;
    brpc::StreamOptions request_stream_options;
    request_stream_options.max_buf_size = 0;
    ASSERT_EQ(0, StreamCreate(&request_stream, cntl, &request_stream_options));
    brpc::ScopedStream stream_guard(request_stream);
    test::EchoService_Stub stub(&channel);
    stub.Echo(&cntl, &request, &response, NULL);
    ASSERT_FALSE(cntl.Failed()) << cntl.ErrorText() << " request_stream=" << request_stream;
    const int N = 10000;
    for (int i = 0; i < N; ++i) {
        int network = htonl(i);
        butil::IOBuf out;
        out.append(&network, sizeof(network));
        ASSERT_EQ(0, brpc::StreamWrite(request_stream, out)) << "i=" << i;
    }
    ASSERT_EQ(0, brpc::StreamClose(request_stream));
    server.Stop(0);
    server.Join();
    while (!handler.stopped()) {
        usleep(100);
    }
    ASSERT_FALSE(handler.failed());
    ASSERT_EQ(0, handler.idle_times());
    ASSERT_EQ(N, handler._expected_next_value);
}

void on_writable(brpc::StreamId, void* arg, int error_code) {
    std::pair<bool, int>* p = (std::pair<bool, int>*)arg;
    p->first = true;
    p->second = error_code;
    LOG(INFO) << "error_code=" << error_code;
}

TEST_F(StreamingRpcTest, block) {
    HandlerControl hc;
    hc.block = true;
    OrderedInputHandler handler(&hc);
    brpc::StreamOptions opt;
    opt.handler = &handler;
    const int N = 10000;
    opt.max_buf_size = sizeof(uint32_t) *  N;
    brpc::Server server;
    MyServiceWithStream service(opt);
    ASSERT_EQ(0, server.AddService(&service, brpc::SERVER_DOESNT_OWN_SERVICE));
    ASSERT_EQ(0, server.Start(9007, NULL));
    brpc::Channel channel;
    ASSERT_EQ(0, channel.Init("127.0.0.1:9007", NULL));
    brpc::Controller cntl;
    brpc::StreamId request_stream;
    brpc::ScopedStream stream_guard(request_stream);
    brpc::StreamOptions request_stream_options;
    request_stream_options.max_buf_size = sizeof(uint32_t) * N;
    ASSERT_EQ(0, StreamCreate(&request_stream, cntl, &request_stream_options));
    test::EchoService_Stub stub(&channel);
    stub.Echo(&cntl, &request, &response, NULL);
    ASSERT_FALSE(cntl.Failed()) << cntl.ErrorText() << " request_stream="
                                << request_stream;
    for (int i = 0; i < N; ++i) {
        int network = htonl(i);
        butil::IOBuf out;
        out.append(&network, sizeof(network));
        ASSERT_EQ(0, brpc::StreamWrite(request_stream, out)) << "i=" << i;
    }
    // sync wait
    int dummy = 102030123;
    butil::IOBuf out;
    out.append(&dummy, sizeof(dummy));
    ASSERT_EQ(EAGAIN, brpc::StreamWrite(request_stream, out));
    hc.block = false;
    // wait flushing all the pending messages
    while (handler._expected_next_value != N) {
        usleep(100);
    }
    // block hanlder again to test async wait
    hc.block = true;
    // async wait
    for (int i = N; i < N + N; ++i) {
        ASSERT_EQ(0, brpc::StreamWait(request_stream, NULL));
        int network = htonl(i);
        butil::IOBuf out;
        out.append(&network, sizeof(network));
        ASSERT_EQ(0, brpc::StreamWrite(request_stream, out)) << "i=" << i;
    }
    out.clear();
    out.append(&dummy, sizeof(dummy));
    ASSERT_EQ(EAGAIN, brpc::StreamWrite(request_stream, out));
    hc.block = false;
    std::pair<bool, int> p = std::make_pair(false, 0);
    usleep(10);
    brpc::StreamWait(request_stream, NULL, on_writable, &p);
    while (!p.first) {
        usleep(100);
    }
    ASSERT_EQ(0, p.second);

    // wait flushing all the pending messages
    while (handler._expected_next_value != N + N) {
        usleep(100);
    }
    usleep(1000);

    LOG(INFO) << "Starting block";
    hc.block = true;
    for (int i = N + N; i < N + N + N; ++i) {
        int network = htonl(i);
        butil::IOBuf out;
        out.append(&network, sizeof(network));
        ASSERT_EQ(0, brpc::StreamWrite(request_stream, out)) << "i=" << i - N - N;
    }
    out.clear();
    out.append(&dummy, sizeof(dummy));
    ASSERT_EQ(EAGAIN, brpc::StreamWrite(request_stream, out));
    timespec duetime = butil::microseconds_from_now(1);
    p.first = false;
    LOG(INFO) << "Start wait";
    brpc::StreamWait(request_stream, &duetime, on_writable, &p);
    while (!p.first) {
        usleep(100);
    }
    ASSERT_TRUE(p.first);
    EXPECT_EQ(ETIMEDOUT, p.second);
    hc.block = false;
    ASSERT_EQ(0, brpc::StreamClose(request_stream));
    while (!handler.stopped()) {
        usleep(100);
    }

    ASSERT_FALSE(handler.failed());
    ASSERT_EQ(0, handler.idle_times());
    ASSERT_EQ(N + N + N, handler._expected_next_value);
}

TEST_F(StreamingRpcTest, auto_close_if_host_socket_closed) {
    HandlerControl hc;
    hc.block = true;
    OrderedInputHandler handler(&hc);
    brpc::StreamOptions opt;
    opt.handler = &handler;
    const int N = 10000;
    opt.max_buf_size = sizeof(uint32_t) *  N;
    brpc::Server server;
    MyServiceWithStream service(opt);
    ASSERT_EQ(0, server.AddService(&service, brpc::SERVER_DOESNT_OWN_SERVICE));
    ASSERT_EQ(0, server.Start(9007, NULL));
    brpc::Channel channel;
    ASSERT_EQ(0, channel.Init("127.0.0.1:9007", NULL));
    brpc::Controller cntl;
    brpc::StreamId request_stream;
    brpc::StreamOptions request_stream_options;
    request_stream_options.max_buf_size = sizeof(uint32_t) * N;
    ASSERT_EQ(0, StreamCreate(&request_stream, cntl, &request_stream_options));
    brpc::ScopedStream stream_guard(request_stream);
    test::EchoService_Stub stub(&channel);
    stub.Echo(&cntl, &request, &response, NULL);
    ASSERT_FALSE(cntl.Failed()) << cntl.ErrorText() << " request_stream=" << request_stream;

    {
        brpc::SocketUniquePtr ptr;
        ASSERT_EQ(0, brpc::Socket::Address(request_stream, &ptr));
        brpc::Stream* s = (brpc::Stream*)ptr->conn();
        ASSERT_TRUE(s->_host_socket != NULL);
        s->_host_socket->SetFailed();
    }

    usleep(100);
    butil::IOBuf out;
    out.append("test");
    ASSERT_EQ(EINVAL, brpc::StreamWrite(request_stream, out));
    while (!handler.stopped()) {
        usleep(100);
    }
    ASSERT_TRUE(handler.failed());
    ASSERT_EQ(0, handler.idle_times());
    ASSERT_EQ(0, handler._expected_next_value);
}

TEST_F(StreamingRpcTest, failed_when_rst) {
    OrderedInputHandler handler;
    brpc::StreamOptions opt;
    opt.handler = &handler;
    opt.messages_in_batch = 100;
    brpc::Server server;
    MyServiceWithStream service(opt);
    ASSERT_EQ(0, server.AddService(&service, brpc::SERVER_DOESNT_OWN_SERVICE));
    ASSERT_EQ(0, server.Start(9007, NULL));
    brpc::Channel channel;
    ASSERT_EQ(0, channel.Init("127.0.0.1:9007", NULL));
    brpc::Controller cntl;
    brpc::StreamId request_stream;
    brpc::StreamOptions request_stream_options;
    request_stream_options.max_buf_size = 0;
    ASSERT_EQ(0, StreamCreate(&request_stream, cntl, &request_stream_options));
    brpc::ScopedStream stream_guard(request_stream);
    test::EchoService_Stub stub(&channel);
    stub.Echo(&cntl, &request, &response, NULL);
    ASSERT_FALSE(cntl.Failed()) << cntl.ErrorText() << " request_stream=" << request_stream;
    const int N = 10000;
    for (int i = 0; i < N; ++i) {
        int network = htonl(i);
        butil::IOBuf out;
        out.append(&network, sizeof(network));
        ASSERT_EQ(0, brpc::StreamWrite(request_stream, out)) << "i=" << i;
    }

    while (handler._expected_next_value != N) {
        usleep(100);
    }
    {
        brpc::SocketUniquePtr ptr;
        ASSERT_EQ(0, brpc::Socket::Address(request_stream, &ptr));
        brpc::Stream* s = (brpc::Stream*)ptr->conn();
        ASSERT_TRUE(s->_host_socket != NULL);
        brpc::policy::SendStreamRst(s->_host_socket,
                                    s->_remote_settings.stream_id());
    }
    // ASSERT_EQ(0, brpc::StreamClose(request_stream));
    server.Stop(0);
    server.Join();
    while (!handler.stopped() && !handler.failed()) {
        usleep(100);
    }
    ASSERT_TRUE(handler.failed());
    ASSERT_EQ(0, handler.idle_times());
    ASSERT_EQ(N, handler._expected_next_value);
}

TEST_F(StreamingRpcTest, idle_timeout) {
    HandlerControl hc;
    hc.block = true;
    OrderedInputHandler handler(&hc);
    brpc::StreamOptions opt;
    opt.handler = &handler;
    opt.idle_timeout_ms = 2;
    const int N = 10000;
    opt.max_buf_size = sizeof(uint32_t) *  N;
    brpc::Server server;
    MyServiceWithStream service(opt);
    ASSERT_EQ(0, server.AddService(&service, brpc::SERVER_DOESNT_OWN_SERVICE));
    ASSERT_EQ(0, server.Start(9007, NULL));
    brpc::Channel channel;
    ASSERT_EQ(0, channel.Init("127.0.0.1:9007", NULL));
    brpc::Controller cntl;
    brpc::StreamId request_stream;
    brpc::StreamOptions request_stream_options;
    request_stream_options.max_buf_size = sizeof(uint32_t) * N;
    ASSERT_EQ(0, StreamCreate(&request_stream, cntl, &request_stream_options));
    brpc::ScopedStream stream_guard(request_stream);
    test::EchoService_Stub stub(&channel);
    stub.Echo(&cntl, &request, &response, NULL);
    ASSERT_FALSE(cntl.Failed()) << cntl.ErrorText() << " request_stream=" << request_stream;
    usleep(10 * 1000 + 800);
    ASSERT_EQ(0, brpc::StreamClose(request_stream));
    while (!handler.stopped()) {
        usleep(100);
    }
    ASSERT_FALSE(handler.failed());
//    ASSERT_TRUE(handler.idle_times() >= 4 && handler.idle_times() <= 6)
//               << handler.idle_times();
    ASSERT_EQ(0, handler._expected_next_value);
}

class PingPongHandler : public brpc::StreamInputHandler {
public:
    int on_received_messages(brpc::StreamId id,
                             butil::IOBuf *const messages[],
                             size_t size) override {
        if (size != 1) {
            LOG(INFO) << "size=" << size;
            _error = true;
            return 0;
        }
        for (size_t i = 0; i < size; ++i) {
            CHECK(messages[i]->length() == sizeof(int));
            int network = 0;
            messages[i]->cutn(&network, sizeof(int));
            if ((int)ntohl(network) != _expected_next_value) {
                _error = true;
            }
            int send_back = ntohl(network) + 1;
            _expected_next_value = send_back + 1;
            butil::IOBuf out;
            network = htonl(send_back);
            out.append(&network, sizeof(network));
            // don't care the return value
            brpc::StreamWrite(id, out);
        }
        return 0;
    }

    void on_idle_timeout(brpc::StreamId /*id*/) override {
        ++_idle_times;
    }

    void on_closed(brpc::StreamId /*id*/) override {
        ASSERT_FALSE(_stopped);
        _stopped = true;
    }


    void on_failed(brpc::StreamId id, int error_code,
                   const std::string& /*error_text*/) override {
        ASSERT_FALSE(_failed);
        ASSERT_NE(0, error_code);
        _failed = true;
    }

    bool error() const { return _error; }
    bool failed() const { return _failed; }
    bool stopped() const { return _stopped; }
    int idle_times() const { return _idle_times; }
private:
    int _expected_next_value{0};
    bool _error{false};
    bool _failed{false};
    bool _stopped{false};
    int _idle_times{0};
};

TEST_F(StreamingRpcTest, ping_pong) {
    PingPongHandler resh;
    brpc::StreamOptions opt;
    opt.handler = &resh;
    const int N = 10000;
    opt.max_buf_size = sizeof(uint32_t) *  N;
    brpc::Server server;
    MyServiceWithStream service(opt);
    ASSERT_EQ(0, server.AddService(&service, brpc::SERVER_DOESNT_OWN_SERVICE));
    ASSERT_EQ(0, server.Start(9007, NULL));
    brpc::Channel channel;
    ASSERT_EQ(0, channel.Init("127.0.0.1:9007", NULL));
    brpc::Controller cntl;
    brpc::StreamId request_stream;
    brpc::StreamOptions request_stream_options;
    PingPongHandler reqh;
    reqh._expected_next_value = 1;
    request_stream_options.handler = &reqh;
    request_stream_options.max_buf_size = sizeof(uint32_t) * N;
    ASSERT_EQ(0, StreamCreate(&request_stream, cntl, &request_stream_options));
    brpc::ScopedStream stream_guard(request_stream);
    test::EchoService_Stub stub(&channel);
    stub.Echo(&cntl, &request, &response, NULL);
    ASSERT_FALSE(cntl.Failed()) << cntl.ErrorText() << " request_stream=" << request_stream;
    int send = 0;
    butil::IOBuf out;
    out.append(&send, sizeof(send));
    ASSERT_EQ(0, brpc::StreamWrite(request_stream, out));
    usleep(10 * 1000);
    ASSERT_EQ(0, brpc::StreamClose(request_stream));
    while (!resh.stopped() || !reqh.stopped()) {
        usleep(100);
    }
    ASSERT_FALSE(resh.error());
    ASSERT_FALSE(reqh.error());
    ASSERT_EQ(0, resh.idle_times());
    ASSERT_EQ(0, reqh.idle_times());
}

class SendNAfterAcceptStream : public AfterAcceptStream {
public:
    explicit SendNAfterAcceptStream(int n)
        : _n(n) {}
    void action(brpc::StreamId s) {
        for (int i = 0; i < _n; ++i) {
            int network = htonl(i);
            butil::IOBuf out;
            out.append(&network, sizeof(network));
            ASSERT_EQ(0, brpc::StreamWrite(s, out)) << "i=" << i;
        }
    }
private:
    int _n;
};

TEST_F(StreamingRpcTest, server_send_data_before_run_done) {
    const int N = 10000;
    SendNAfterAcceptStream after_accept(N);
    brpc::StreamOptions opt;
    opt.max_buf_size = -1;
    brpc::Server server;
    MyServiceWithStream service(opt, &after_accept);
    ASSERT_EQ(0, server.AddService(&service, brpc::SERVER_DOESNT_OWN_SERVICE));
    ASSERT_EQ(0, server.Start(9007, NULL));
    brpc::Channel channel;
    ASSERT_EQ(0, channel.Init("127.0.0.1:9007", NULL));
    OrderedInputHandler handler;
    brpc::StreamOptions request_stream_options;
    request_stream_options.handler = &handler;
    brpc::StreamId request_stream;
    brpc::Controller cntl;
    ASSERT_EQ(0, StreamCreate(&request_stream, cntl, &request_stream_options));
    brpc::ScopedStream stream_guard(request_stream);
    test::EchoService_Stub stub(&channel);
    stub.Echo(&cntl, &request, &response, NULL);
    ASSERT_FALSE(cntl.Failed()) << cntl.ErrorText() << " request_stream=" << request_stream;
    // wait flushing all the pending messages
    while (handler._expected_next_value != N) {
        usleep(100);
    }
    ASSERT_EQ(0, brpc::StreamClose(request_stream));
    while (!handler.stopped()) {
        usleep(100);
    }
    ASSERT_FALSE(handler.failed());
    ASSERT_EQ(0, handler.idle_times());
}

TEST_F(StreamingRpcTest, segment_stream_data_automatically) {
    GFLAGS_NAMESPACE::SetCommandLineOption("stream_write_max_segment_size", "1");
    OrderedInputHandler handler;
    brpc::StreamOptions opt;
    opt.handler = &handler;
    opt.messages_in_batch = 100;
    brpc::Server server;
    MyServiceWithStream service(opt);
    ASSERT_EQ(0, server.AddService(&service, brpc::SERVER_DOESNT_OWN_SERVICE));
    ASSERT_EQ(0, server.Start(9007, NULL));
    brpc::Channel channel;
    ASSERT_EQ(0, channel.Init("127.0.0.1:9007", NULL));
    brpc::Controller cntl;
    brpc::StreamId request_stream;
    brpc::StreamOptions request_stream_options;
    ASSERT_EQ(0, StreamCreate(&request_stream, cntl, &request_stream_options));
    brpc::ScopedStream stream_guard(request_stream);
    test::EchoService_Stub stub(&channel);
    stub.Echo(&cntl, &request, &response, NULL);
    ASSERT_FALSE(cntl.Failed()) << cntl.ErrorText() << " request_stream=" << request_stream;
    const int N = 1000;
    for (int i = 0; i < N; ++i) {
        int network = htonl(i);
        butil::IOBuf out;
        out.append(&network, sizeof(network));
        ASSERT_EQ(0, brpc::StreamWrite(request_stream, out)) << "i=" << i;
    }

    brpc::SocketUniquePtr host_socket_ptr;
    {
      brpc::SocketUniquePtr ptr;
      ASSERT_EQ(0, brpc::Socket::Address(request_stream, &ptr));
      brpc::Stream *s = (brpc::Stream *)ptr->conn();
      ASSERT_TRUE(s->_host_socket != NULL);
      s->_host_socket->ReAddress(&host_socket_ptr);
    }

    ASSERT_EQ(0, brpc::StreamClose(request_stream));
    server.Stop(0);
    server.Join();
    while (!handler.stopped()) {
        usleep(100);
    }
    const int64_t now_ms = butil::cpuwide_time_ms();
    host_socket_ptr->UpdateStatsEverySecond(now_ms);
    brpc::SocketStat stat;
    host_socket_ptr->GetStat(&stat);
    ASSERT_LT(N * sizeof(N), stat.out_num_messages_m);
    ASSERT_FALSE(handler.failed());
    ASSERT_EQ(0, handler.idle_times());
    ASSERT_EQ(N, handler._expected_next_value);
    GFLAGS_NAMESPACE::SetCommandLineOption("stream_write_max_segment_size", "536870912");
}