adbd.c

来自「基于DHT的对等协议」· C语言 代码 · 共 1,538 行 · 第 1/3 页

	if (rt < exptime) {
	  str filepath = subdirpath << "/" << sdp->d_name;
	  if (unlink (filepath) < 0)
	    warn ("unlink: %s: %m\n", filepath.cstr ());
	}
      }
      closedir (subdir);
      rmdir (subdirpath);
    }
  }
  closedir (datadir);
  return 0;
}
// }}}
// }}}

// {{{ DB Manager
class dbmanager {
  str dbpath;
  str logpath;
  ihash<str, dbns, &dbns::name, &dbns::hlink> dbs;
  void dbcloser (dbns *db);
public:
  dbmanager (str, str);
  ~dbmanager ();
  const str &getdbpath () { return dbpath; };

  dbns *get (const str &n) { return dbs[n]; };
  dbns *createdb (const str &n, bool aux);
};

void
mkdir_wrapper (str path)
{
  // Make path if necessary
  struct stat sb;
  if (stat (path, &sb) < 0) {
    if (errno == ENOENT) {
      if (mkdir (path, 0755) < 0)
	fatal ("mkdir (%s): %m\n", path.cstr ());
    } else {
      fatal ("stat (%s): %m\n", path.cstr ());
    }
  } else {
    if (!S_ISDIR (sb.st_mode)) 
      fatal ("%s is not a directory\n", path.cstr ());
    if (access (path, W_OK) < 0)
      fatal ("access (%s, W_OK): %m\n", path.cstr ());
  }
}

static str
canonicalize_path (str path)
{
  // Convert path to full path
  char realpath[MAXPATHLEN];
  int fd = -1;
  if (((fd = open (".", O_RDONLY)) >= 0)
     && chdir (path) >= 0)
  {
    if (getcwd (realpath, sizeof (realpath)))
      errno = 0;
    else if (!errno)
      errno = EINVAL;
    if (fchdir (fd))
      warn ("fchdir: %m\n");
  }
  return str (realpath);
}

dbmanager::dbmanager (str p, str lp = NULL) :
  dbpath (p),
  logpath (lp)
{
  mkdir_wrapper (dbpath);
  if (logpath) {
    mkdir_wrapper (logpath);
    logpath = canonicalize_path (logpath);
  }

  // Add trailing slash if necessary
  str x = canonicalize_path (dbpath);
  if (x[x.len () - 1] != '/')
    dbpath = strbuf () << x << "/";
  else 
    dbpath = x;
}

dbmanager::~dbmanager ()
{
  dbs.traverse (wrap (this, &dbmanager::dbcloser));
  dbs.clear ();
}

void
dbmanager::dbcloser (dbns *db)
{
  if (db) {
    warn << "Closing db " << db->name << "\n";
    delete db;
    db = NULL;
  }
}

dbns *
dbmanager::createdb (const str &n, bool aux)
{
  dbns *db = dbs[n];
  if (!db) {
    if (logpath)
      mkdir_wrapper (strbuf() << logpath << "/" << n);
    db = New dbns (dbpath, n, aux, logpath);
    dbs.insert (db);
  }
  return db;
}
// }}}

// {{{ Shutdown functions
EXITFN (cleanup);
static void
cleanup ()
{
  if (dbstarted) {
    if (clntfd >= 0) {
      fdcb (clntfd, selread, NULL);
      close (clntfd);
    }
    unlink (dbsock);
  }
}

static void
halt ()
{
  warn << "Exiting on command...\n";
  delete dbm;
  exit (0);
}
// }}}

// {{{ RPC execution
// {{{ do_initspace
void
do_initspace (dbmanager *dbm, svccb *sbp)
{
  adb_initspacearg *arg = sbp->Xtmpl getarg<adb_initspacearg> ();
  adb_status stat (ADB_OK);
  dbns *db = dbm->get (arg->name);
  if (db) {
    sbp->replyref (stat); 
    return;
  }
  db = dbm->createdb (arg->name, arg->hasaux);
  stat = (db ? ADB_OK : ADB_ERR);
  sbp->replyref (stat);
}
// }}}
// {{{ do_store
void
do_store (dbmanager *dbm, svccb *sbp)
{

  adb_storearg *arg = sbp->Xtmpl getarg<adb_storearg> ();
  dbns *db = dbm->get (arg->name);
  if (!db) {
    sbp->replyref (ADB_ERR);
    return;
  }

  u_int64_t t = io_start ();

  DBT data;
  bzero (&data, sizeof (data));
  data.data = arg->data.base ();
  data.size = arg->data.size ();

  int r = db->insert (arg->key, data, arg->auxdata, arg->expiration);
  io_finish (t, strbuf("store %s", arg->name.cstr ()));
  adb_status stat = ADB_OK;
  switch (r) {
    case 0:
    case DB_KEYEXIST:
      break;
    case ENOSPC:
      stat = ADB_DISKFULL;
      break;
    default:
      stat = ADB_ERR;
      break;
  }
  sbp->replyref (stat);

  if (db->quotacheck (quota) > expire_threshold) {
    u_int64_t t = io_start ();
    db->expire (expire_batch_size);
    io_finish (t, strbuf ("expire %s", db->getname ().cstr ()));
  }
}
// }}}
// {{{ do_fetch
void
do_fetch (dbmanager *dbm, svccb *sbp)
{
  adb_fetcharg *arg = sbp->Xtmpl getarg<adb_fetcharg> ();
  adb_fetchres res (ADB_OK); 

  dbns *db = dbm->get (arg->name);
  if (!db) {
    res.set_status (ADB_ERR);
    sbp->replyref (res);
    return;
  }

  u_int64_t t = io_start ();
 
  str data; 
  adb_metadata_t md;
  bzero (&md, sizeof (md));

  int r;
  chordID key;
  if (arg->nextkey) {
    r = db->lookup_nextkey (arg->key, key);
    data = "";
  } else {
    r = db->lookup (arg->key, data, md);
    key = arg->key;
  }

  if (r) {
    res.set_status ((r == DB_NOTFOUND ? ADB_NOTFOUND : ADB_ERR));
  } else {
    res.resok->key = key;
    res.resok->data = data;
    res.resok->expiration = md.expiration;
  }

  io_finish (t, strbuf ("fetch %s", arg->name.cstr ()));

  sbp->replyref (res);
}
// }}}
// {{{ do_getkeys
void
do_getkeys (dbmanager *dbm, svccb *sbp)
{
  adb_getkeysarg *arg = sbp->Xtmpl getarg<adb_getkeysarg> ();
  adb_getkeysres res (ADB_OK);

  dbns *db = dbm->get (arg->name);
  if (!db) {
    res.set_status (ADB_ERR);
    sbp->replyref (res);
    return;
  }
  res.resok->hasaux = db->hasaux () && arg->getaux;
  res.resok->ordered = arg->ordered;

  int r (-1);
  r = db->getkeys (arg->continuation, arg->batchsize, arg->getaux, res.resok->keyaux);
  if (!r)
    res.resok->continuation = incID (res.resok->keyaux.back ().key);
  res.resok->complete = (r == DB_NOTFOUND);
  if (r && r != DB_NOTFOUND) 
    res.set_status (ADB_ERR);
  sbp->replyref (res);
}
// }}}
// {{{ do_delete
void
do_delete (dbmanager *dbm, svccb *sbp)
{
  adb_deletearg *arg = sbp->Xtmpl getarg<adb_deletearg> ();
  dbns *db = dbm->get (arg->name);
  if (!db) {
    sbp->replyref (ADB_ERR);
    return;
  }
  int r = db->del (arg->key, arg->auxdata);
  sbp->replyref ((r == 0) ? ADB_OK : ADB_NOTFOUND);
}

// }}}
// {{{ do_getspaceinfo
void
do_getspaceinfo (dbmanager *dbm, svccb *sbp)
{
  adb_dbnamearg *arg = sbp->Xtmpl getarg<adb_dbnamearg> ();
  adb_getspaceinfores *res = sbp->Xtmpl getres<adb_getspaceinfores> ();
  res->status = ADB_OK;
  dbns *db = dbm->get (arg->name);
  if (!db) {
    sbp->replyref (ADB_ERR);
    return;
  }

  res->fullpath = strbuf () << dbm->getdbpath () << arg->name;
  res->hasaux = db->hasaux ();
  sbp->reply (res);
}
// }}}
// {{{ do_sync
void
do_sync (dbmanager *dbm, svccb *sbp)
{
  adb_dbnamearg *arg = sbp->Xtmpl getarg<adb_dbnamearg> ();
  adb_status res (ADB_OK);
  dbns *db = dbm->get (arg->name);
  if (!db) {
    res = ADB_ERR;
  } else {
    u_int64_t t = io_start ();
    db->sync ();
    io_finish (t, strbuf ("sync %s", arg->name.cstr ()));
  }
  sbp->replyref (res);
}
// }}}
// {{{ do_expire
void
do_expire (dbmanager *dbm, svccb *sbp)
{
  adb_expirearg *arg = sbp->Xtmpl getarg<adb_expirearg> ();
  adb_status res (ADB_OK);
  dbns *db = dbm->get (arg->name);
  if (!db) {
    res = ADB_ERR;
  } else {
    u_int64_t t = io_start ();
    db->expire (arg->limit, arg->deadline);
    io_finish (t, strbuf ("expire %s", arg->name.cstr ()));
  }
  sbp->replyref (res);
}
// }}}
// }}}

// {{{ RPC accept and dispatch
void
dispatch (ref<axprt_stream> s, ptr<asrv> a, dbmanager *dbm, svccb *sbp)
{
  if (sbp == NULL) {
    warn << "EOF from client\n";
    a->setcb (NULL);
    return;
  }

  switch (sbp->proc ()) {
  case ADBPROC_INITSPACE:
    do_initspace (dbm, sbp);
    break;
  case ADBPROC_STORE:
    do_store (dbm, sbp);
    break;
  case ADBPROC_FETCH:
    do_fetch (dbm, sbp);
    break;
  case ADBPROC_GETKEYS:
    do_getkeys (dbm, sbp);
    break;
  case ADBPROC_DELETE:
    do_delete (dbm, sbp);
    break;
  case ADBPROC_GETSPACEINFO:
    do_getspaceinfo (dbm, sbp);
    break;
  case ADBPROC_SYNC:
    do_sync (dbm, sbp);
    break;
  case ADBPROC_EXPIRE:
    do_expire (dbm, sbp);
    break;
  default:
    fatal << "unknown procedure: " << sbp->proc () << "\n";
  }
  
  return;
}

static void accept_cb (int lfd, dbmanager *dbm);
static void
listen_unix (str sock_name, dbmanager *dbm)
{
  unlink (sock_name);
  clntfd = unixsocket (sock_name);
  if (clntfd < 0) 
    fatal << "Error creating socket (UNIX)" << strerror (errno) << "\n";
  
  if (listen (clntfd, 128) < 0) {
    fatal ("Error from listen: %m\n");
    close (clntfd);
  } else {
    fdcb (clntfd, selread, wrap (accept_cb, clntfd, dbm));
  }
  dbstarted = true;
}

static void 
accept_cb (int lfd, dbmanager *dbm)
{
  sockaddr_un sun;
  bzero (&sun, sizeof (sun));
  socklen_t sunlen = sizeof (sun);
  int fd = accept (lfd, reinterpret_cast<sockaddr *> (&sun), &sunlen);
  if (fd < 0)
    fatal ("EOF\n");

  ref<axprt_stream> x = axprt_stream::alloc (fd, asrvbufsize);

  ptr<asrv> a = asrv::alloc (x, adb_program_1);
  a->setcb (wrap (dispatch, x, a, dbm));
}
// }}}

void
usage ()
{
  warnx << "Usage: adbd -d db -S sock [-D] [-q quota]\n";
  exit (0);
}

int 
main (int argc, char **argv)
{

  setprogname (argv[0]);
  mp_set_memory_functions (NULL, simple_realloc, NULL);

  char ch;
  str db_name = "/var/tmp/db";
  str log_path = NULL;
  dbsock = "/tmp/db-sock";

  bool do_daemonize (false);

  while ((ch = getopt (argc, argv, "Dd:l:q:S:"))!=-1)
    switch (ch) {
    case 'D':
      do_daemonize = true;
      break;
    case 'd':
      db_name = optarg;
      break;
    case 'l':
      log_path = optarg;
      break;
    case 'q':
      {
	u_int64_t factor = u_int64_t (1024) * 1024 * 1024;
	char f = optarg[strlen (optarg) - 1];
	if (!isdigit (f))
	  optarg[strlen (optarg) - 1] = 0;
	bool ok = convertint (optarg, &quota);
	if (!ok)
	  usage ();
	switch (f) {
	  case 'b':
	    factor = 1;
	    break;
	  case 'K':
	    factor = 1024;
	    break;
	  case 'M':
	    factor = 1024 * 1024;
	    break;
	  case 'G':
	    factor = u_int64_t (1024) * 1024 * 1024;
	    break;
	  default:
	    if (!isdigit (f))
	      fatal ("Unknown conversion factor '%c'\n", f);
	}
	quota *= factor;
      }
      break;
    case 'S':
      dbsock = optarg;
      break;
    default:
      usage ();
    }
 
  if (do_daemonize) {
    warn << "starting daemonized\n";
    daemonize ();
  }
  warn << "db path: " << db_name << "\n";
  if (log_path)
    warn << "log path: " << log_path << "\n";
  warn << "db socket: " << dbsock << "\n";

  dbm = New dbmanager (db_name, log_path);

  sigcb (SIGHUP, wrap (&halt));
  sigcb (SIGINT, wrap (&halt));
  sigcb (SIGTERM, wrap (&halt));

  sigcb (SIGUSR1, wrap (&toggle_iotime));

  bzero (&master_metadata, sizeof (master_metadata));
  master_metadata.data = strdup ("MASTER_INFO");
  master_metadata.size = strlen ("MASTER_INFO");

  //setup the asrv
  listen_unix (dbsock, dbm);

  amain ();
}

// vim: foldmethod=marker