mabain/binaries/mbc.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 // A mabain command-line client
  20
  21 #include <stdlib.h>
  22 #include <string.h>
  23 #include <iostream>
  24 #include <fstream>
  25 #include <string>
  26 #include <assert.h>
  27 #include <signal.h>
  28 #include <readline/readline.h>
  29 #include <readline/history.h>
  30
  31 #include "db.h"
  32 #include "mb_data.h"
  33 #include "dict.h"
  34 #include "error.h"
  35 #include "version.h"
  36
  37 #include "hexbin.h"
  38 #include "expr_parser.h"
  39
  40 using namespace mabain;
  41
  42 enum mbc_command {
  43     COMMAND_NONE = 0,
  44     COMMAND_QUIT = 1,
  45     COMMAND_UNKNOWN = 2,
  46     COMMAND_STATS = 3,
  47     COMMAND_FIND = 4,
  48     COMMAND_FIND_ALL = 5,
  49     COMMAND_INSERT = 6,
  50     COMMAND_REPLACE = 7,
  51     COMMAND_DELETE = 8,
  52     COMMAND_DELETE_ALL = 9,
  53     COMMAND_HELP = 10,
  54     COMMAND_RESET_N_WRITER = 11,
  55     COMMAND_RESET_N_READER = 12,
  56     COMMAND_FIND_LPREFIX = 13,
  57     COMMAND_PRINT_HEADER = 14,
  58     COMMAND_FIND_HEX = 15,
  59     COMMAND_FIND_LPREFIX_HEX = 16,
  60     COMMAND_RECLAIM_RESOURCES = 17,
  61     COMMAND_PARSING_ERROR = 18,
  62 };
  63
  64 volatile bool quit_mbc = false;
  65 static void HandleSignal(int sig)
  66 {
  67     switch(sig)
  68     {
  69         case SIGSEGV:
  70             std::cerr << "process segfault\n";
  71             abort();
  72         case SIGTERM:
  73         case SIGINT:
  74         case SIGQUIT:
  75         case SIGHUP:
  76         case SIGPIPE:
  77             quit_mbc = true;
  78         case SIGUSR1:
  79         case SIGUSR2:
  80             break;
  81     }
  82 }
  83
  84 static void usage(const char *prog)
  85 {
  86     std::cout << "Usage: " << prog << " -d mabain-directory [-im index-memcap] [-dm data-memcap] [-w] [-e query] [-s script-file]\n";
  87     std::cout <<"\t-d mabain databse directory\n";
  88     std::cout <<"\t-im index memcap\n";
  89     std::cout <<"\t-dm data memcap\n";
  90     std::cout <<"\t-w running in writer mode\n";
  91     std::cout <<"\t-e run query on command line\n";
  92     std::cout <<"\t-s run queries in a file\n";
  93     exit(1);
  94 }
  95
  96 static void show_help()
  97 {
  98     std::cout << "\tfind(\"key\")\t\tsearch entry by key\n";
  99     std::cout << "\tfindPrefix(\"key\")\tsearch entry by key using longest prefix match\n";
 100     std::cout << "\tfindAll\t\t\tlist all entries\n";
 101     std::cout << "\tinsert(\"key\":\"value\")\tinsert a key-value pair\n";
 102     std::cout << "\treplace(\"key\":\"value\")\treplace a key-value pair\n";
 103     std::cout << "\tdelete(\"key\")\t\tdelete entry by key\n";
 104     std::cout << "\tdeleteAll\t\tdelete all entries\n";
 105     std::cout << "\tshow\t\t\tshow database statistics\n";
 106     std::cout << "\thelp\t\t\tshow helps\n";
 107     std::cout << "\tquit\t\t\tquit mabain client\n";
 108     std::cout << "\tdecWriterCount\t\tClear writer count in shared memory header\n";
 109     std::cout << "\tdecReaderCount\t\tdecrement reader count in shared memory header\n";
 110     std::cout << "\tprintHeader\t\tPrint shared memory header\n";
 111     std::cout << "\treclaimResources\tReclaim deleted resources\n";
 112 }
 113
 114 static void trim_spaces(const char *cmd, std::string &cmd_trim)
 115 {
 116     cmd_trim.clear();
 117
 118     int quotation = 0;
 119     const char *p = cmd;
 120     while(*p != '\0')
 121     {
 122         if(*p == '\'' || *p == '\"')
 123         {
 124             cmd_trim.append(1, '\'');
 125             quotation ^= 1;
 126         }
 127         else if(!isspace(*p) || quotation)
 128         {
 129             cmd_trim.append(1, *p);
 130         }
 131
 132         p++;
 133     }
 134 }
 135
 136 static bool check_hex_output(std::string &cmd)
 137 {
 138     size_t pos = cmd.rfind(".hex()");
 139     if(pos == std::string::npos)
 140         return false;
 141     if(pos == cmd.length()-6)
 142     {
 143         cmd.erase(pos);
 144         return true;
 145     }
 146
 147     return false;
 148 }
 149
 150 static int parse_key_value_pair(const std::string &kv_str,
 151                                 std::string &key,
 152                                 std::string &value)
 153 {
 154     // search for ':' that separates key and value pair
 155     // currently this utility does not support quotation in
 156     // quotation, e.g, "abc\"def" as key or value.
 157     size_t pos = 0;
 158     int quotation_cnt = 0;
 159     for(size_t i = 0; i < kv_str.length(); i++)
 160     {
 161         if(kv_str[i] == '\'')
 162         {
 163             quotation_cnt++;
 164         }
 165         else if(kv_str[i] == ':')
 166         {
 167             // do not count the ':' in the key or value string.
 168             if(quotation_cnt % 2 == 0)
 169             {
 170                 pos = i;
 171                 break;
 172             }
 173         }
 174     }
 175
 176     if(pos == 0)
 177         return -1;
 178
 179     ExprParser expr_key(kv_str.substr(0, pos));
 180     if(expr_key.Evaluate(key) < 0)
 181         return -1;
 182
 183     ExprParser expr_value(kv_str.substr(pos+1));
 184     if(expr_value.Evaluate(value) < 0)
 185         return -1;
 186
 187     return 0;
 188 }
 189
 190 static int parse_command(std::string &cmd,
 191                          std::string &key,
 192                          std::string &value)
 193 {
 194     std::string yes;
 195     bool hex_output = false;
 196
 197     key = "";
 198     value = "";
 199     switch(cmd[0])
 200     {
 201         case 'q':
 202             if(cmd.compare("quit") == 0)
 203                 return COMMAND_QUIT;
 204             break;
 205         case 's':
 206             if(cmd.compare("show") == 0)
 207                 return COMMAND_STATS;
 208             break;
 209         case 'f':
 210             hex_output = check_hex_output(cmd);
 211             if(cmd.compare(0, 5, "find(") == 0)
 212             {
 213                 if(cmd[cmd.length()-1] != ')')
 214                     return COMMAND_UNKNOWN;
 215                 ExprParser expr(cmd.substr(5, cmd.length()-6));
 216                 if(expr.Evaluate(key) < 0)
 217                     return COMMAND_PARSING_ERROR;
 218                 if(hex_output)
 219                     return COMMAND_FIND_HEX;
 220                 return COMMAND_FIND;
 221             }
 222             else if(cmd.compare(0, 11, "findPrefix(") == 0)
 223             {
 224                 if(cmd[cmd.length()-1] != ')')
 225                     return COMMAND_UNKNOWN;
 226                 ExprParser expr(cmd.substr(11, cmd.length()-12));
 227                 if(expr.Evaluate(key) < 0)
 228                     return COMMAND_PARSING_ERROR;
 229                 if(hex_output)
 230                     return COMMAND_FIND_LPREFIX_HEX;
 231                 return COMMAND_FIND_LPREFIX;
 232             }
 233             else if(cmd.compare("findAll") == 0)
 234                 return COMMAND_FIND_ALL;
 235             break;
 236         case 'd':
 237             if(cmd.compare(0, 7, "delete(") == 0)
 238             {
 239                 if(cmd[cmd.length()-1] != ')')
 240                     return COMMAND_UNKNOWN;
 241                 ExprParser expr(cmd.substr(7, cmd.length()-8));
 242                 if(expr.Evaluate(key) < 0)
 243                     return COMMAND_PARSING_ERROR;
 244                 return COMMAND_DELETE;
 245             }
 246             else if(cmd.compare("deleteAll") == 0)
 247             {
 248                 std::cout << "Do you want to delete all entries? Press \'Y\' to continue: ";
 249                 std::string del_all;
 250                 std::getline(std::cin, del_all);
 251                 if(del_all.length() == 0 || del_all[0] != 'Y')
 252                     return COMMAND_NONE;
 253                 return COMMAND_DELETE_ALL;
 254             }
 255             else if(cmd.compare("decReaderCount") == 0)
 256             {
 257                 std::cout << "Do you want to decrement number of reader? Press \'y\' to continue: ";
 258                 std::getline(std::cin, yes);
 259                 if(yes.length() > 0 && yes[0] == 'y')
 260                     return COMMAND_RESET_N_READER;
 261                 return COMMAND_NONE;
 262             }
 263             else if(cmd.compare("decWriterCount") == 0)
 264             {
 265                 std::cout << "Do you want to decrement number of writer? Press \'y\' to continue: ";
 266                 std::getline(std::cin, yes);
 267                 if(yes.length() > 0 && yes[0] == 'y')
 268                     return COMMAND_RESET_N_WRITER;
 269                 return COMMAND_NONE;
 270             }
 271             break;
 272         case 'i':
 273             if(cmd.compare(0, 7, "insert(") == 0)
 274             {
 275                 if(cmd[cmd.length()-1] != ')')
 276                     return COMMAND_UNKNOWN;
 277                 if(parse_key_value_pair(cmd.substr(7, cmd.length()-8), key, value) < 0)
 278                     return COMMAND_PARSING_ERROR;
 279                 return COMMAND_INSERT;
 280             }
 281             break;
 282         case 'r':
 283             if(cmd.compare(0, 8, "replace(") == 0)
 284             {
 285                 if(cmd[cmd.length()-1] != ')')
 286                     return COMMAND_UNKNOWN;
 287                 if(parse_key_value_pair(cmd.substr(8, cmd.length()-9), key, value) < 0)
 288                     return COMMAND_PARSING_ERROR;
 289                 return COMMAND_REPLACE;
 290             }
 291             else if(cmd.compare("reclaimResources") == 0)
 292                 return COMMAND_RECLAIM_RESOURCES;
 293             else
 294                 return COMMAND_UNKNOWN;
 295             break;
 296         case 'h':
 297             if(cmd.compare("help") == 0)
 298                 return COMMAND_HELP;
 299             break;
 300         case 'p':
 301             if(cmd.compare("printHeader") == 0)
 302                 return COMMAND_PRINT_HEADER;
 303             break;
 304         default:
 305             break;
 306     }
 307
 308     return COMMAND_UNKNOWN;
 309 }
 310
 311 #define ENTRY_PER_PAGE 20
 312 static void display_all_kvs(DB *db)
 313 {
 314     if(db == NULL)
 315         return;
 316
 317     int count = 0;
 318     for(DB::iterator iter = db->begin(); iter != db->end(); ++iter)
 319     {
 320         count++;
 321         std::cout << iter.key << ": " <<
 322                      std::string((char *)iter.value.buff, iter.value.data_len) << "\n";
 323         if(count % ENTRY_PER_PAGE == 0)
 324         {
 325             std::string show_more;
 326             std::cout << "Press \'y\' for displaying more: ";
 327             std::getline(std::cin, show_more);
 328             if(show_more.length() == 0 || show_more[0] != 'y')
 329                 break;
 330         }
 331     }
 332 }
 333
 334 static void display_output(const MBData &mbd, bool hex_output, bool prefix)
 335 {
 336     if(prefix)
 337         std::cout << "key length matched: " << mbd.match_len << "\n";
 338     if(hex_output)
 339     {
 340         char hex_buff[256];
 341         int len = mbd.data_len;
 342         if(256 < 2*len + 1)
 343         {
 344             std::cout << "display the first 127 bytes\n";
 345             len = 127;
 346         }
 347         if(bin_2_hex((const uint8_t *)mbd.buff, len, hex_buff, 256) < 0)
 348             std::cout << "failed to convert binary to hex\n";
 349         else
 350             std::cout << hex_buff << "\n";
 351     }
 352     else
 353     {
 354         std::cout << std::string((char *)mbd.buff, mbd.data_len) << "\n";
 355     }
 356 }
 357
 358 static int RunCommand(int mode, DB *db, int cmd_id, const std::string &key, const std::string &value)
 359 {
 360     int rval = MBError::SUCCESS;
 361     bool overwrite = false;
 362     bool hex_output = false;
 363     MBData mbd;
 364
 365     switch(cmd_id)
 366     {
 367         case COMMAND_NONE:
 368             // no opertation needed
 369             break;
 370         case COMMAND_QUIT:
 371             std::cout << "bye\n";
 372             quit_mbc = true;
 373             break;
 374         case COMMAND_FIND_HEX:
 375             hex_output = true;
 376         case COMMAND_FIND:
 377             rval = db->Find(key, mbd);
 378             if(rval == MBError::SUCCESS)
 379                 display_output(mbd, hex_output, false);
 380             else
 381                 std::cout << MBError::get_error_str(rval) << "\n";
 382             break;
 383         case COMMAND_FIND_LPREFIX_HEX:
 384             hex_output = true;
 385         case COMMAND_FIND_LPREFIX:
 386             rval = db->FindLongestPrefix(key, mbd);
 387             if(rval == MBError::SUCCESS)
 388                 display_output(mbd, hex_output, true);
 389             else
 390                 std::cout << MBError::get_error_str(rval) << "\n";
 391             break;
 392         case COMMAND_DELETE:
 393             if(mode & CONSTS::ACCESS_MODE_WRITER)
 394             {
 395                 rval = db->Remove(key);
 396                 std::cout << MBError::get_error_str(rval) << "\n";
 397             }
 398             else
 399                 std::cout << "permission not allowed\n";
 400             break;
 401         case COMMAND_REPLACE:
 402             overwrite = true;
 403         case COMMAND_INSERT:
 404             if(mode & CONSTS::ACCESS_MODE_WRITER)
 405             {
 406                 rval = db->Add(key, value, overwrite);
 407                 std::cout << MBError::get_error_str(rval) << "\n";
 408             }
 409             else
 410                 std::cout << "permission not allowed\n";
 411             break;
 412         case COMMAND_STATS:
 413             db->PrintStats();
 414             break;
 415         case COMMAND_HELP:
 416             show_help();
 417             break;
 418         case COMMAND_DELETE_ALL:
 419             if(mode & CONSTS::ACCESS_MODE_WRITER)
 420             {
 421                 rval = db->RemoveAll();
 422                 std::cout << MBError::get_error_str(rval) << "\n";
 423             }
 424             else
 425                 std::cout << "permission not allowed\n";
 426             break;
 427         case COMMAND_FIND_ALL:
 428             display_all_kvs(db);
 429             break;
 430         case COMMAND_RESET_N_WRITER:
 431             if(mode & CONSTS::ACCESS_MODE_WRITER)
 432                 std::cout << "writer is running, cannot reset writer counter\n";
 433             else
 434                 db->UpdateNumHandlers(CONSTS::ACCESS_MODE_WRITER, -1);
 435             break;
 436         case COMMAND_RESET_N_READER:
 437             db->UpdateNumHandlers(CONSTS::ACCESS_MODE_READER, -1);
 438             break;
 439         case COMMAND_PRINT_HEADER:
 440             db->PrintHeader();
 441             break;
 442         case COMMAND_RECLAIM_RESOURCES:
 443             if(mode & CONSTS::ACCESS_MODE_WRITER)
 444                 db->CollectResource(1, 1);
 445             else
 446                 std::cout << "writer is not running, can not perform grabage collection" << std::endl;
 447             break;
 448         case COMMAND_PARSING_ERROR:
 449             break;
 450         case COMMAND_UNKNOWN:
 451         default:
 452             std::cout << "unknown query\n";
 453             break;
 454     }
 455
 456     return rval;
 457 }
 458
 459 static void mbclient(DB *db, int mode)
 460 {
 461     rl_bind_key('\t', rl_complete);
 462
 463     printf("mabain %d.%d.%d shell\n", version[0], version[1], version[2]);
 464     std::cout << "database directory: " << db->GetDBDir() << "\n";
 465
 466     int cmd_id;
 467     std::string key;
 468     std::string value;
 469     std::string cmd;
 470
 471     while(true)
 472     {
 473         char* line = readline(">> ");
 474         if(line == NULL) break;
 475         if(line[0] == '\0')
 476         {
 477             free(line);
 478             continue;
 479         }
 480
 481         trim_spaces(line, cmd);
 482         add_history(line);
 483         free(line);
 484         cmd_id = parse_command(cmd, key, value);
 485
 486         RunCommand(mode, db, cmd_id, key, value);
 487
 488         if(quit_mbc) break;
 489     }
 490 }
 491
 492 static void run_query_command(DB *db, int mode, const std::string &command_str)
 493 {
 494     std::string cmd;
 495     int cmd_id;
 496     std::string key;
 497     std::string value;
 498
 499     trim_spaces(command_str.c_str(), cmd);
 500     if(cmd.length() == 0)
 501     {
 502         std::cerr << command_str << " not a valid command\n";
 503         return;
 504     }
 505
 506     cmd_id = parse_command(cmd, key, value);
 507     RunCommand(mode, db, cmd_id, key, value);
 508 }
 509
 510 static void run_script(DB *db, int mode, const std::string &script_file)
 511 {
 512     std::ifstream script_in(script_file);
 513     if(!script_in.is_open()) {
 514         std::cerr << "cannot open file " << script_file << "\n";
 515         return;
 516     }
 517
 518     std::string line;
 519     int cmd_id;
 520     std::string key;
 521     std::string value;
 522     std::string cmd;
 523
 524     while(getline(script_in, line))
 525     {
 526         trim_spaces(line.c_str(), cmd);
 527         if(cmd.length() == 0)
 528         {
 529             std::cerr << line << " not a valid query\n";
 530             continue;
 531         }
 532
 533         cmd_id = parse_command(cmd, key, value);
 534         std::cout << cmd << ": ";
 535         RunCommand(mode, db, cmd_id, key, value);
 536
 537         if(quit_mbc) break;
 538     }
 539     script_in.close();
 540 }
 541
 542 int main(int argc, char *argv[])
 543 {
 544     sigset_t mask;
 545
 546     signal(SIGINT, HandleSignal);
 547     signal(SIGTERM, HandleSignal);
 548     signal(SIGQUIT, HandleSignal);
 549     signal(SIGPIPE, HandleSignal);
 550     signal(SIGHUP, HandleSignal);
 551     signal(SIGSEGV, HandleSignal);
 552     signal(SIGUSR1, HandleSignal);
 553     signal(SIGUSR2, HandleSignal);
 554
 555     sigemptyset(&mask);
 556     sigaddset(&mask, SIGTERM);
 557     sigaddset(&mask, SIGQUIT);
 558     sigaddset(&mask, SIGINT);
 559     sigaddset(&mask, SIGPIPE);
 560     sigaddset(&mask, SIGHUP);
 561     sigaddset(&mask, SIGSEGV);
 562     sigaddset(&mask, SIGUSR1);
 563     sigaddset(&mask, SIGUSR2);
 564     pthread_sigmask(SIG_SETMASK, &mask, NULL);
 565
 566     sigemptyset(&mask);
 567     pthread_sigmask(SIG_SETMASK, &mask, NULL);
 568
 569     int64_t memcap_i = 1024*1024LL;
 570     int64_t memcap_d = 1024*1024LL;
 571     const char *db_dir = NULL;
 572     int mode = 0;
 573     std::string query_cmd = "";
 574     std::string script_file = "";
 575     int64_t index_blk_size = 64LL*1024*1024;
 576     int64_t data_blk_size = 64LL*1024*1024;
 577     int64_t lru_bucket_size = 1000;
 578
 579     for(int i = 1; i < argc; i++)
 580     {
 581         if(strcmp(argv[i], "-d") == 0)
 582         {
 583             if(++i >= argc)
 584                 usage(argv[0]);
 585             db_dir = argv[i];
 586         }
 587         else if(strcmp(argv[i], "-im") == 0)
 588         {
 589             if(++i >= argc)
 590                 usage(argv[0]);
 591             memcap_i = atoi(argv[i]);
 592         }
 593         else if(strcmp(argv[i], "-dm") == 0)
 594         {
 595             if(++i >= argc)
 596                 usage(argv[0]);
 597             memcap_d = atoi(argv[i]);
 598         }
 599         else if(strcmp(argv[i], "-w") == 0)
 600         {
 601             mode |= CONSTS::ACCESS_MODE_WRITER;
 602         }
 603         else if(strcmp(argv[i], "-e") == 0)
 604         {
 605             if(++i >= argc)
 606                 usage(argv[0]);
 607             query_cmd = argv[i];
 608         }
 609         else if(strcmp(argv[i], "-s") == 0)
 610         {
 611             if(++i >= argc)
 612                 usage(argv[0]);
 613             script_file = argv[i];
 614         }
 615         else if(strcmp(argv[i], "--lru-bucket-size") == 0)
 616         {
 617             if(++i >= argc)
 618                 usage(argv[0]);
 619             lru_bucket_size = atoi(argv[i]);
 620         }
 621         else if(strcmp(argv[i], "--index-block-size") == 0)
 622         {
 623             if(++i >= argc)
 624                 usage(argv[0]);
 625             index_blk_size = atoi(argv[i]);
 626         }
 627         else if(strcmp(argv[i], "--data-block-size") == 0)
 628         {
 629             if(++i >= argc)
 630                 usage(argv[0]);
 631             data_blk_size = atoi(argv[i]);
 632         }
 633         else
 634             usage(argv[0]);
 635     }
 636
 637     if(db_dir == NULL)
 638         usage(argv[0]);
 639
 640     MBConfig mbconf;
 641     memset(&mbconf, 0, sizeof(mbconf));
 642     mbconf.mbdir = db_dir;
 643     mbconf.options = mode;
 644     mbconf.memcap_index = memcap_i;
 645     mbconf.memcap_data = memcap_d;
 646     mbconf.block_size_index = index_blk_size;
 647     mbconf.block_size_data = data_blk_size;
 648     mbconf.num_entry_per_bucket = lru_bucket_size;
 649     DB *db = new DB(mbconf);
 650     if(!db->is_open())
 651     {
 652         std::cout << db->StatusStr() << "\n";
 653         exit(1);
 654     }
 655
 656     // DB::SetLogFile("/var/tmp/mabain.log");
 657     // DB::LogDebug();
 658
 659     if(query_cmd.length() != 0)
 660     {
 661         run_query_command(db, mode, query_cmd);
 662     }
 663     else if(script_file.length() != 0)
 664     {
 665         run_script(db, mode, script_file);
 666     }
 667     else
 668     {
 669         mbclient(db, mode);
 670     }
 671
 672     db->Close();
 673     // DB::CloseLogFile();
 674     delete db;
 675     return 0;
 676 }