mabain/examples/mb_multi_proc_test.cpp

   1 /**
   2  * Copyright (C) 2017 Cisco Inc.
   3  *
   4  * This program is free software: you can redistribute it and/or  modify
   5  * it under the terms of the GNU General Public License, version 2,
   6  * as published by the Free Software Foundation.
   7  *
   8  * This program is distributed in the hope that it will be useful,
   9  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  10  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  11  * GNU General Public License for more details.
  12  *
  13  * You should have received a copy of the GNU General Public License
  14  * along with this program.  If not, see <http://www.gnu.org/licenses/>.
  15  */
  16
  17 // @author Changxue Deng <chadeng@cisco.com>
  18
  19 // Multiple writers/readers access the DB concurrently.
  20
  21 #include <unistd.h>
  22 #include <sys/wait.h>
  23 #include <string.h>
  24 #include <assert.h>
  25 #include <sys/time.h>
  26 #include <iostream>
  27 #include <string>
  28 #include <vector>
  29 #include <openssl/sha.h>
  30 #include <map>
  31
  32 #include <mabain/db.h>
  33
  34 #define ADD_TEST             0
  35 #define REMOVE_ONE_BY_ONE    1
  36 #define REMOVE_ALL           2
  37 #define SHRINK_TEST          3
  38 #define LOOKUP_TEST          4
  39 #define ITERATOR_TEST        5
  40 #define KEY_TYPE_INT         0
  41 #define KEY_TYPE_SHA256      1
  42
  43 const char* mb_dir = "/var/tmp/mabain_test";
  44
  45 using namespace mabain;
  46
  47 static int stop_processing = 0;
  48 static int num_keys = 1000000;
  49 static int test_type = 0;
  50 static int64_t memcap_index = 256*1024*1024LL;
  51 static int64_t memcap_data  = 256*1024*1024LL;
  52 static int key_type = KEY_TYPE_SHA256;
  53 //static int key_type = KEY_TYPE_INT;
  54 static std::map<std::string,int> checked;
  55
  56 static const char* get_sha256_str(int key);
  57
  58 static void SetTestStatus(bool success, bool stop_test = true)
  59 {
  60     std::string cmd;
  61     if(success) {
  62         cmd = std::string("touch ") + mb_dir + "/_success";
  63     } else {
  64         cmd = std::string("rm ") + mb_dir + "/_success";
  65     }
  66     if(system(cmd.c_str()) != 0) {
  67     }
  68     if(stop_test) {
  69         abort();
  70     }
  71 }
  72
  73 static void RemoveRandom(int rm_cnt, DB &db)
  74 {
  75     int cnt = 0;
  76     std::string kv;
  77     int ikey = 0;
  78     while(cnt <= rm_cnt) {
  79         cnt++;
  80
  81         ikey = rand() % num_keys + 1;
  82         switch(key_type) {
  83             case KEY_TYPE_SHA256:
  84                 kv = get_sha256_str(ikey);
  85                 break;
  86             case KEY_TYPE_INT:
  87             default:
  88                 kv = std::string("key") + std::to_string(ikey);
  89                 break;
  90         }
  91
  92         db.Remove(kv);
  93         checked[kv] = 0;
  94         if(cnt % 1000000 == 0)
  95             std::cout << "Removed " << cnt << "\n";
  96     }
  97 }
  98
  99 static void PopulateDB(DB &db)
 100 {
 101     int rval;
 102     int nkeys = 0;
 103     std::string kv;
 104     while(nkeys < num_keys) {
 105         nkeys++;
 106
 107         switch(key_type) {
 108             case KEY_TYPE_SHA256:
 109                 kv = get_sha256_str(nkeys);
 110                 break;
 111             case KEY_TYPE_INT:
 112             default:
 113                 kv = std::string("key") + std::to_string(nkeys);
 114                 break;
 115         }
 116
 117         rval = db.Add(kv, kv);
 118         if(rval != MBError::SUCCESS)
 119         {
 120             std::cout << "failed to add " << kv << " " << rval << "\n";
 121             SetTestStatus(false);
 122         }
 123         checked[kv] = 1;
 124
 125         if(nkeys % 1000000 == 0)
 126             std::cout << "populate db:  " << " inserted " << nkeys << " keys\n";
 127     }
 128 }
 129
 130 static void RemoveHalf(DB &db)
 131 {
 132     int rval;
 133     int nkeys = 0;
 134     std::string kv;
 135     while(nkeys < num_keys) {
 136         nkeys++;
 137         if(nkeys % 1000000 == 0)
 138             std::cout << "Removed " << nkeys/2 << "\n";
 139         if(nkeys % 2 == 0)
 140             continue;
 141
 142         switch(key_type) {
 143             case KEY_TYPE_SHA256:
 144                 kv = get_sha256_str(nkeys);
 145                 break;
 146             case KEY_TYPE_INT:
 147             default:
 148                 kv = std::string("key") + std::to_string(nkeys);
 149                 break;
 150         }
 151         rval = db.Remove(kv.c_str(), kv.length());
 152         checked[kv] = 0;
 153         if(rval != MBError::SUCCESS)
 154         {
 155             std::cout << "failed to remove " << nkeys << " " << kv << " " << rval << "\n";
 156             SetTestStatus(false);
 157         }
 158     }
 159 }
 160
 161 //Writer process
 162 static void Writer(int id)
 163 {
 164     DB db(mb_dir, CONSTS::WriterOptions(), memcap_index, memcap_data);
 165     if(!db.is_open()) {
 166         std::cerr << "writer " << id << " failed to open mabain db: "
 167                   << db.StatusStr() << "\n";
 168         exit(1);
 169     }
 170
 171     if(test_type == SHRINK_TEST) {
 172         db.CollectResource(0, 0);
 173
 174         db.Close();
 175         exit(0);
 176     }
 177
 178     int rval;
 179     int nkeys = 0;
 180     std::string kv;
 181     int count = 0;
 182     while(!stop_processing) {
 183         nkeys++;
 184
 185         switch(key_type) {
 186             case KEY_TYPE_SHA256:
 187                 kv = get_sha256_str(nkeys);
 188                 break;
 189             case KEY_TYPE_INT:
 190             default:
 191                 kv = std::string("key") + std::to_string(nkeys);
 192                 break;
 193         }
 194
 195         if(test_type == REMOVE_ONE_BY_ONE) {
 196             rval = db.Remove(kv.c_str(), kv.length());
 197             if(rval != MBError::SUCCESS && rval != MBError::NOT_EXIST) {
 198                 std::cout << "failed to remove " << kv << "\n";
 199                 SetTestStatus(false);
 200             } else {
 201                 count++;
 202             }
 203         } else if(test_type == REMOVE_ALL) {
 204             if(nkeys == int(num_keys * 0.666666)) {
 205                 db.RemoveAll();
 206             }
 207         } else {
 208             rval = db.Add(kv, kv);
 209             if(rval != MBError::SUCCESS && rval != MBError::IN_DICT) {
 210                 std::cout << "failed to add " << kv << " " << rval << "\n";
 211                 SetTestStatus(false);
 212             } else if(rval == MBError::SUCCESS) {
 213                 count++;
 214             }
 215         }
 216
 217         if(nkeys % 1000000 == 0)
 218             std::cout << "writer " << id << " looped over " << nkeys << " keys\n";
 219         if(nkeys >= num_keys) break;
 220     }
 221
 222     db.Close();
 223     exit(0);
 224 }
 225
 226 //Reader process
 227 static void Reader(int id)
 228 {
 229     DB db(mb_dir, CONSTS::ReaderOptions(), memcap_index, memcap_data);
 230     if(!db.is_open()) {
 231         std::cerr << "reader " << id << " failed to open mabain db: "
 232                   << db.StatusStr() << "\n";
 233         SetTestStatus(false);
 234     }
 235
 236     std::string key;
 237     if(test_type == ITERATOR_TEST) {
 238         int count = 0;
 239         for(DB::iterator iter = db.begin(); iter != db.end(); ++iter) {
 240             count++;
 241             if(iter.key != std::string((char*)iter.value.buff, iter.value.data_len)) {
 242                 std::cout << "VALUE NOT MATCH " << iter.key << ": "
 243                           << std::string((char*)iter.value.buff, iter.value.data_len) << "\n";
 244                 SetTestStatus(false);
 245             } else {
 246                 if(checked[iter.key] == 1)
 247                     checked[iter.key]++;
 248             }
 249         }
 250         db.Close();
 251         for(int i = 1; i <= num_keys; i++) {
 252             switch(key_type) {
 253                 case KEY_TYPE_SHA256:
 254                     key = get_sha256_str(i);
 255                     break;
 256                 case KEY_TYPE_INT:
 257                 default:
 258                     key = std::string("key") + std::to_string(i);
 259                     break;
 260             }
 261             if(!(checked[key] == 0 || checked[key] == 2)) {
 262                 std::cerr << i << ": " << key << " " << checked[key] << " for id " << id << "\n";
 263                 SetTestStatus(false);
 264             }
 265         }
 266         exit(0);
 267     }
 268
 269     int ikey = 1;
 270     MBData mb_data;
 271     int rval;
 272     int nfound = 0;
 273     while(!stop_processing) {
 274         switch(key_type) {
 275             case KEY_TYPE_SHA256:
 276                 key = get_sha256_str(ikey);
 277                 break;
 278             case KEY_TYPE_INT:
 279             default:
 280                 key = std::string("key") + std::to_string(ikey);
 281                 break;
 282         }
 283         rval = db.Find(key, mb_data);
 284
 285         if(rval == MBError::SUCCESS) {
 286             // check the value read from DB
 287             if(memcmp(key.c_str(), mb_data.buff, mb_data.data_len)) {
 288                 std::cout << "READER " << id << " VALUE DOES NOT MATCH: "
 289                           << key << ":" << mb_data.match_len << " "
 290                           << std::string((const char *)mb_data.buff, mb_data.data_len) << "\n";
 291                 SetTestStatus(false);
 292             }
 293
 294             ikey++;
 295             nfound++;
 296         }
 297         else if(rval == MBError::NOT_EXIST) {
 298             ikey++;
 299         } else {
 300             std::cerr << "ERROR: " << MBError::get_error_str(rval) << "\n";
 301             SetTestStatus(false);
 302         }
 303
 304         if(ikey % 1000000 == 0)
 305             std::cout << "reader " << id << " looked up " << ikey << " keys\n";
 306
 307         if(ikey > num_keys) break;
 308     }
 309
 310     stop_processing = 1;
 311     db.Close();
 312     exit(0);
 313 }
 314
 315 static char sha256_str[65];
 316 static const char* get_sha256_str(int key)
 317 {
 318     unsigned char hash[SHA256_DIGEST_LENGTH];
 319     SHA256_CTX sha256;
 320     SHA256_Init(&sha256);
 321     SHA256_Update(&sha256, (unsigned char*)&key, 4);
 322     SHA256_Final(hash, &sha256);
 323     int i = 0;
 324     for(i = 0; i < SHA256_DIGEST_LENGTH; i++)
 325     {
 326         sprintf(sha256_str + (i * 2), "%02x", hash[i]);
 327     }
 328     sha256_str[64] = 0;
 329     return (const char*)sha256_str;
 330 }
 331
 332 int main(int argc, char *argv[])
 333 {
 334     int num_writer = 1; // There can be one writer at most.
 335     int num_reader = 2;
 336     test_type = -1;
 337     srand(time(NULL));
 338
 339     mabain::DB::SetLogFile("/var/tmp/mabain_test/mabain.log");
 340     if(argc >= 4) {
 341         const char *test_tp = argv[1];
 342         if(strcmp(test_tp, "add") == 0)
 343             test_type = ADD_TEST;
 344         else if(strcmp(test_tp, "remove") == 0)
 345             test_type = REMOVE_ONE_BY_ONE;
 346         else if(strcmp(test_tp, "remove-all") == 0)
 347             test_type = REMOVE_ALL;
 348         else if(strcmp(test_tp, "shrink") == 0)
 349             test_type = SHRINK_TEST;
 350         else if(strcmp(test_tp, "lookup") == 0)
 351             test_type = LOOKUP_TEST;
 352         else if(strcmp(test_tp, "iterator") == 0)
 353             test_type = ITERATOR_TEST;
 354         num_keys = atoi(argv[2]);
 355         num_reader = atoi(argv[3]);
 356     }
 357     if(argc >= 5) {
 358         mb_dir = argv[4];
 359     }
 360
 361     SetTestStatus(false, false);
 362     if(test_type < 0) {
 363         test_type = ADD_TEST;
 364     }
 365     // Delete all keys from last run.
 366     DB db(mb_dir, CONSTS::WriterOptions(), memcap_index, memcap_data);
 367     if(!db.is_open()) {
 368         std::cerr << " failed to open mabain db: "
 369                   << db.StatusStr() << "\n";
 370         exit(1);
 371     }
 372
 373     db.RemoveAll();
 374     PopulateDB(db);
 375     switch(test_type) {
 376         case ADD_TEST:
 377             RemoveRandom(int(num_keys * 0.7777777), db);
 378             break;
 379         case LOOKUP_TEST:
 380             break;
 381         case SHRINK_TEST:
 382             RemoveHalf(db);
 383             break;
 384         case REMOVE_ONE_BY_ONE:
 385             break;
 386         case REMOVE_ALL:
 387             break;
 388         case ITERATOR_TEST:
 389             RemoveRandom(int(num_keys * 0.666666), db);
 390             break;
 391         default:
 392             abort();
 393     }
 394
 395     db.Close();
 396
 397     std::vector<pid_t> pid_arr;
 398     pid_t test_pid = getpid();
 399
 400     if(test_type == LOOKUP_TEST)
 401         num_writer = 0;
 402     assert(num_writer < 2);
 403     struct timeval start_tm;
 404     gettimeofday(&start_tm,NULL);
 405     for(int i = 0; i < num_writer; i++) {
 406         if(getpid() != test_pid) continue;
 407
 408         pid_t pid = fork();
 409         if(pid < 0) break;
 410
 411         if(getpid() == test_pid) {
 412             pid_arr.push_back(pid);
 413             continue;
 414         }
 415
 416         // Start child/writer process
 417         Writer(i);
 418     }
 419
 420     for(int i = 0; i < num_reader; i++) {
 421         if(getpid() != test_pid) continue;
 422
 423         pid_t pid = fork();
 424         if(pid < 0) break;
 425
 426         if(getpid() == test_pid) {
 427             pid_arr.push_back(pid);
 428             continue;
 429         }
 430
 431         // Start child/reader process
 432         Reader(i);
 433     }
 434
 435     if(test_pid == getpid()) {
 436         int status;
 437         int i;
 438         while(!stop_processing) {
 439             // Check if all children exited
 440             for(i = 0; i < (int) pid_arr.size(); i++) {
 441                 if(pid_arr[i] > 0) {
 442                     waitpid(pid_arr[i], &status, WUNTRACED | WCONTINUED);
 443                     pid_arr[i] = -1;
 444
 445                     if(i == (int) pid_arr.size()-1) {
 446                         stop_processing = 1;
 447                         struct timeval stop_tm;
 448                         gettimeofday(&stop_tm,NULL);
 449                         //std::cout << "All children have exited!\n";
 450                         //std::cout << "time: " << ((stop_tm.tv_sec-start_tm.tv_sec)*1000000.0 +
 451                         //            (stop_tm.tv_usec-start_tm.tv_usec))/num_keys << "\n";
 452                     }
 453                 }
 454             }
 455
 456             usleep(1000);
 457         }
 458     }
 459
 460     mabain::DB::CloseLogFile();
 461     SetTestStatus(true, false);
 462     return 0;
 463 }