brpc_socket_map_unittest.cpp 4.9 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17
// 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.

18
// brpc - A framework to host and access services throughout Baidu.
gejun's avatar
gejun committed
19 20 21 22

// Date: Sun Jul 13 15:04:18 CST 2014

#include <gtest/gtest.h>
23
#include <gflags/gflags.h>
gejun's avatar
gejun committed
24 25 26 27 28 29 30 31 32 33 34
#include "brpc/socket.h"
#include "brpc/socket_map.h"
#include "brpc/reloadable_flags.h"

namespace brpc {
DECLARE_int32(idle_timeout_second);
DECLARE_int32(defer_close_second);
DECLARE_int32(max_connection_pool_size);
} // namespace brpc

namespace {
35
butil::EndPoint g_endpoint;
36
brpc::SocketMapKey g_key(g_endpoint);
gejun's avatar
gejun committed
37 38 39 40 41 42 43 44

void* worker(void*) {
    const int ROUND = 2;
    const int COUNT = 1000;
    brpc::SocketId id;
    for (int i = 0; i < ROUND * 2; ++i) {
        for (int j = 0; j < COUNT; ++j) {
            if (i % 2 == 0) {
45
                EXPECT_EQ(0, brpc::SocketMapInsert(g_key, &id));
gejun's avatar
gejun committed
46
            } else {
47
                brpc::SocketMapRemove(g_key);
gejun's avatar
gejun committed
48 49 50 51 52 53 54 55 56 57 58 59 60 61 62
            }
        }
    }
    return NULL;
}

class SocketMapTest : public ::testing::Test{
protected:
    SocketMapTest(){};
    virtual ~SocketMapTest(){};
    virtual void SetUp(){};
    virtual void TearDown(){};
};

TEST_F(SocketMapTest, idle_timeout) {
63
    const int TIMEOUT = 1;
gejun's avatar
gejun committed
64 65 66 67 68 69 70 71 72 73 74
    const int NTHREAD = 10;
    brpc::FLAGS_defer_close_second = TIMEOUT;
    pthread_t tids[NTHREAD];
    for (int i = 0; i < NTHREAD; ++i) {
        ASSERT_EQ(0, pthread_create(&tids[i], NULL, worker, NULL));
    }
    for (int i = 0; i < NTHREAD; ++i) {
        ASSERT_EQ(0, pthread_join(tids[i], NULL));
    }
    brpc::SocketId id;
    // Socket still exists since it has not reached timeout yet
75
    ASSERT_EQ(0, brpc::SocketMapFind(g_key, &id));
76
    usleep(TIMEOUT * 1000000L + 1100000L);
gejun's avatar
gejun committed
77
    // Socket should be removed after timeout
78
    ASSERT_EQ(-1, brpc::SocketMapFind(g_key, &id));
gejun's avatar
gejun committed
79 80

    brpc::FLAGS_defer_close_second = TIMEOUT * 10;
81 82 83
    ASSERT_EQ(0, brpc::SocketMapInsert(g_key, &id));
    brpc::SocketMapRemove(g_key);
    ASSERT_EQ(0, brpc::SocketMapFind(g_key, &id));
gejun's avatar
gejun committed
84 85
    // Change `FLAGS_idle_timeout_second' to 0 to disable checking
    brpc::FLAGS_defer_close_second = 0;
86
    usleep(1100000L);
gejun's avatar
gejun committed
87
    // And then Socket should be removed
88
    ASSERT_EQ(-1, brpc::SocketMapFind(g_key, &id));
gejun's avatar
gejun committed
89 90

    brpc::SocketId main_id;
91
    ASSERT_EQ(0, brpc::SocketMapInsert(g_key, &main_id));
gejun's avatar
gejun committed
92 93 94 95
    brpc::FLAGS_idle_timeout_second = TIMEOUT;
    brpc::SocketUniquePtr main_ptr;
    brpc::SocketUniquePtr ptr;
    ASSERT_EQ(0, brpc::Socket::Address(main_id, &main_ptr));
gejun's avatar
gejun committed
96
    ASSERT_EQ(0, main_ptr->GetPooledSocket(&ptr));
gejun's avatar
gejun committed
97 98 99 100 101
    ASSERT_TRUE(main_ptr.get());
    main_ptr.reset();
    id = ptr->id();
    ptr->ReturnToPool();
    ptr.reset(NULL);
102
    usleep(TIMEOUT * 1000000L + 2000000L);
gejun's avatar
gejun committed
103 104 105 106
    // Pooled connection should be `ReleaseAdditionalReference',
    // which destroyed the Socket. As a result `GetSocketFromPool'
    // should return a new one
    ASSERT_EQ(0, brpc::Socket::Address(main_id, &main_ptr));
gejun's avatar
gejun committed
107
    ASSERT_EQ(0, main_ptr->GetPooledSocket(&ptr));
gejun's avatar
gejun committed
108 109 110
    ASSERT_TRUE(main_ptr.get());
    main_ptr.reset();
    ASSERT_NE(id, ptr->id());
111
    brpc::SocketMapRemove(g_key);
gejun's avatar
gejun committed
112 113 114 115 116 117 118 119
}

TEST_F(SocketMapTest, max_pool_size) {
    const int MAXSIZE = 5;
    const int TOTALSIZE = MAXSIZE + 5;
    brpc::FLAGS_max_connection_pool_size = MAXSIZE;

    brpc::SocketId main_id;
120
    ASSERT_EQ(0, brpc::SocketMapInsert(g_key, &main_id));
gejun's avatar
gejun committed
121 122 123 124 125

    brpc::SocketUniquePtr ptrs[TOTALSIZE];
    for (int i = 0; i < TOTALSIZE; ++i) {
        brpc::SocketUniquePtr main_ptr;
        ASSERT_EQ(0, brpc::Socket::Address(main_id, &main_ptr));
gejun's avatar
gejun committed
126
        ASSERT_EQ(0, main_ptr->GetPooledSocket(&ptrs[i]));
gejun's avatar
gejun committed
127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143
        ASSERT_TRUE(main_ptr.get());
        main_ptr.reset();
    }
    for (int i = 0; i < TOTALSIZE; ++i) {
        ASSERT_EQ(0, ptrs[i]->ReturnToPool());
    }
    std::vector<brpc::SocketId> ids;
    brpc::SocketUniquePtr main_ptr;
    ASSERT_EQ(0, brpc::Socket::Address(main_id, &main_ptr));
    main_ptr->ListPooledSockets(&ids);
    EXPECT_EQ(MAXSIZE, (int)ids.size());
    // The last few Sockets should be `SetFailed' by `ReturnSocketToPool'
    for (int i = MAXSIZE; i < TOTALSIZE; ++i) {
        EXPECT_TRUE(ptrs[i]->Failed());
    }
}
} //namespace
144 145

int main(int argc, char* argv[]) {
146
    butil::str2endpoint("127.0.0.1:12345", &g_key.peer.addr);
147 148 149
    testing::InitGoogleTest(&argc, argv);
    return RUN_ALL_TESTS();
}