gtest_db_thread.cpp 2.0 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14
// Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed 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.

W
wangjiawei04 已提交
15
#include "kvdb/rocksdb_impl.h"
16 17 18 19 20 21 22 23 24 25 26 27 28 29
#include "kvdb/kvdb_impl.h"
#include "kvdb/paddle_rocksdb.h"
#include <gtest/gtest.h>
#include <string>
#include <fstream>
#include <chrono>
#include <thread>
class KVDBTest : public ::testing::Test {
protected:
    void SetUp() override{
                
    }
    
    static void SetUpTestCase() {
30

31 32 33 34 35 36 37
    }
};
int my_argc;
char** my_argv;


void db_thread_test(AbsKVDBPtr kvdb, int size) {
W
wangjiawei04 已提交
38
    for (int i = 0; i < size; i++) {
39 40 41 42 43 44
        kvdb->Set(std::to_string(i), std::to_string(i));
        kvdb->Get(std::to_string(i));
    }
}

TEST_F(KVDBTest, AbstractKVDB_Thread_Test) {
W
wangjiawei04 已提交
45
    if (my_argc != 3) {
46 47 48 49 50 51
        std::cerr << "illegal input! should be db_thread ${num_of_thread} ${num_of_ops_each_thread}" << std::endl;
        return;
    }
    int num_of_thread = atoi(my_argv[1]);
    int nums_of_ops_each_thread = atoi(my_argv[2]);
    std::vector<AbsKVDBPtr> kvdbptrs;
W
wangjiawei04 已提交
52
    for (int i= 0; i < num_of_thread; i++) {
53 54 55 56
        kvdbptrs.push_back(std::make_shared<RocksKVDB>());
        kvdbptrs[i]->CreateDB();
    }
    std::vector<std::thread> tarr;
W
wangjiawei04 已提交
57
    for (int i = 0; i< num_of_thread; i++) {
58 59
        tarr.push_back(std::thread(db_thread_test, kvdbptrs[i], nums_of_ops_each_thread));
    }
W
wangjiawei04 已提交
60
    for (int i = 0; i< num_of_thread; i++) {
61 62 63 64 65 66 67 68 69 70 71 72 73 74
        tarr[i].join();
    }
    return;
}



int main(int argc, char** argv) {
    my_argc = argc;
    my_argv = argv;
    ::testing::InitGoogleTest(&argc, argv);
     return RUN_ALL_TESTS();
}