merkle_tree_disk.c

基于DHT的对等协议

#include <chord_types.h>
#include <id_utils.h>
#include <rxx.h>
#include "merkle_tree_disk.h"
#include "dhash_common.h"

static FILE *
open_file (str name) {
  // make all the parent directories if applicable
  vec<str> dirs;
  static const rxx dirsplit("\\/");
  (void) split (&dirs, dirsplit, name);
  strbuf filestrbuf;
  for (uint i = 0; i < dirs.size()-1; i++) {
    filestrbuf << dirs[i] << "/";
    
    str dirpath (filestrbuf);
    // check for existence, make if necessary
    struct stat st;
    if (stat (dirpath, &st) < 0) {
      if (mkdir(dirpath, S_IRWXU|S_IRWXG|S_IROTH|S_IXOTH) != 0) {
	fatal << "open_file: mkdir (" << dirpath << "): " << strerror (errno) << ".\n";
      }
    }
  }

  FILE *f = fopen (name, "r+");
  if (f == NULL) {
    f = fopen (name, "w+");
  }
  if (f == NULL) {
    fatal << "open_file: fopen (" << name << "): " << strerror (errno) << ".\n"; 
  }
  return f;
}

static int
copy_file (str src, str dst)
{
  int sfd = open (src, O_RDONLY);
  if (sfd < 0) {
    perror ("copy_file: src");
    return sfd;
  }

  str tmpfile = dst << "~";
  unlink (tmpfile);
  int dfd = open (tmpfile, O_CREAT|O_WRONLY, 0644);
  if (dfd < 0) {
    perror ("copy_file: dst");
    close (sfd);
    return dfd;
  }
#define COPY_ERROR(msg) \
  do {						  \
    int oerrno = errno;				  \
    warn ("copy_file: " msg ": %m");		  \
    close (sfd); close (dfd); unlink (tmpfile);	  \
    errno = oerrno;				  \
    return -1;					  \
  } while (0)

  char buf[8192];
  ssize_t nread;
  while ((nread = read (sfd, buf, sizeof (buf))) > 0) {
    ssize_t nwrote = write (dfd, buf, nread);
    if (nwrote < 0) {
      COPY_ERROR ("bad write");
    }
    if (nwrote != nread) {
      COPY_ERROR ("short write");
    }
  }
  if (nread < 0) {
    COPY_ERROR ("read error");
  }
  close (sfd); close (dfd);
  rename (tmpfile, dst);
#undef COPY_ERROR
  return 0;
}

//////////////// merkle_node_disk /////////////////

merkle_node_disk::merkle_node_disk (FILE *internal, FILE *leaf, 
				    MERKLE_DISK_TYPE type, u_int32_t block_no) :
  merkle_node (),
  hashes (NULL),
  children (NULL),
  _internal (internal),
  _leaf (leaf), 
  _type (type),
  _block_no (block_no)
{
  if (isleaf()) {
    merkle_leaf_node leaf;
    int seekval = fseek (_leaf, _block_no*sizeof (merkle_leaf_node), SEEK_SET);
    assert (seekval == 0);
    fread (&leaf, sizeof (merkle_leaf_node), 1, _leaf);

    count = ntohl(leaf.key_count);
    for (uint i = 0; i < count; i++) {
      chordID c;
      mpz_set_rawmag_be (&c, leaf.keys[i].key, sizeof (leaf.keys[i].key));
      keylist.insert (New merkle_key (c));
    }
  } else {
    children = New array<u_int32_t, 64> ();
    hashes = New array<merkle_hash_id, 64> ();

    merkle_internal_node internal;
    int seekval = fseek (_internal, _block_no*sizeof (merkle_internal_node), 
			 SEEK_SET);
    assert (seekval == 0);
    fread (&internal, sizeof (merkle_internal_node), 1, _internal);

    count = ntohl(internal.key_count);
    for (uint i = 0; i < 64; i++) {
      mpz_set_rawmag_be (&((*hashes)[i].id), internal.hashes[i].key, 
			 sizeof(internal.hashes[i].key));
      (*hashes)[i].hash = merkle_hash ((*hashes)[i].id);
      (*children)[i] = ntohl (internal.child_pointers[i]);
    }
  }

  rehash ();
}

merkle_node_disk::~merkle_node_disk () 
{
  keylist.deleteall_correct();
  delete children;
  delete hashes;

  // also delete all children
  for (uint i = 0; i < to_delete.size(); i++) {
    delete to_delete[i];
  }

  // not necessary, but helps catch dangling pointers
  bzero (this, sizeof (*this)); 
}

void merkle_node_disk::rehash ()
{
  hash = 0;
  if (count == 0) {
    return;
  }
  
  sha1ctx sc;
  if (isleaf ()) {
    assert (count > 0 && count <= 64);
    merkle_key *k = keylist.first();
    while (k != NULL) {
      merkle_hash h (k->id);
      sc.update (h.bytes, h.size);
      k = keylist.next(k);
    }
  } else {
    for (uint i = 0; i < 64; i++) {
      merkle_hash h = (*hashes)[i].hash;
      sc.update (h.bytes, h.size);
    }
  }
  sc.final (hash.bytes);
}

void
merkle_node_disk::write_out ()
{
  if (isleaf ()) {
    merkle_leaf_node leaf;
    bzero (&leaf, sizeof (leaf));
    leaf.key_count = htonl (count);
    merkle_key *m = keylist.first();
    int i = 0;
    while (m != NULL) {
      mpz_get_rawmag_be (leaf.keys[i].key, sizeof(leaf.keys[i].key), &(m->id));
      m = keylist.next (m);
      i++;
    }

    int seekval = fseek (_leaf, _block_no*sizeof(merkle_leaf_node), SEEK_SET);
    assert (seekval == 0);
    fwrite (&leaf, sizeof(merkle_leaf_node), 1, _leaf);
  } else {
    merkle_internal_node internal;
    internal.key_count = htonl(count);
    for (uint i = 0; i < 64; i++) {
      mpz_get_rawmag_be (internal.hashes[i].key, 
			 sizeof (internal.hashes[i].key), &((*hashes)[i].id));
      internal.child_pointers[i] = htonl ((*children)[i]);
      //      warn << _block_no << ") writing out child " << i << ") " << ((*children)[i] >> 1) << "\n";
    }

    int seekval = fseek (_internal, _block_no*sizeof(merkle_internal_node), 
			 SEEK_SET);
    assert (seekval == 0);
    fwrite (&internal, sizeof(merkle_internal_node), 1, _internal);
  }
}

merkle_node *
merkle_node_disk::child (u_int i)
{
  assert (!isleaf ());
  assert (i >= 0 && i < 64);

  u_int32_t pointer = (*children)[i];
  MERKLE_DISK_TYPE type;
  if (pointer % 2 == 0) {
    type = MERKLE_DISK_INTERNAL;
  } else {
    type = MERKLE_DISK_LEAF;
  }
  merkle_node_disk *n = New merkle_node_disk (_internal, _leaf, type, 
					      pointer >> 1);
  to_delete.push_back (n);
  return n;
}

merkle_hash
merkle_node_disk::child_hash (u_int i)
{
  assert (!isleaf ());
  assert (i >= 0 && i < 64);
  return (*hashes)[i].hash;
}

u_int32_t
merkle_node_disk::child_ptr (u_int i)
{
  assert (!isleaf ());
  assert (i >= 0 && i < 64);
  return (*children)[i];
}

void
merkle_node_disk::set_child (merkle_node_disk *n, u_int i)
{
  assert (!isleaf ());
  assert (i >= 0 && i < 64);

  (*children)[i] = ((n->get_block_no() << 1) | (n->isleaf()?0x00000001:0));
  (*hashes)[i].hash = n->hash;
  (*hashes)[i].id = static_cast<bigint> (n->hash);
}

void
merkle_node_disk::add_key (chordID key)
{
  assert (isleaf () && count < 64);
  count++;
  merkle_key *m = New merkle_key (key);
  keylist.insert (m);
}

void
merkle_node_disk::add_key (merkle_hash key)
{
  assert (isleaf () && count < 64);
  count++;
  merkle_key *m = New merkle_key(key);
  keylist.insert(m);
}

void
merkle_node_disk::internal2leaf ()
{
  assert (!isleaf ());
  _type = MERKLE_DISK_LEAF;
  delete children;
  children = NULL;
  delete hashes;
  hashes = NULL;
}

bool
merkle_node_disk::isleaf () const {
  return (_type == MERKLE_DISK_LEAF);
}

void merkle_node_disk::leaf2internal () {
  assert (isleaf ());
  _type = MERKLE_DISK_INTERNAL;
  children = New array<u_int32_t, 64>();
  hashes = New array<merkle_hash_id, 64>();
  keylist.deleteall_correct();
  assert (!isleaf ());
}

static void
indent (u_int depth)
{
  while (depth-- > 0)
    warnx << " ";
}

void
merkle_node_disk::dump (u_int depth)
{
  warnx << "[NODE "
	<< strbuf ("0x%x", (u_int)this)
	<< " cnt:" << count
	<< " hash:" << hash
	<< ">\n";
  err_flush ();

  merkle_node *n = dynamic_cast<merkle_node *> (this);
  if (!n->isleaf ()) {
    for (int i = 0; i < 64; i++) {
      merkle_node *child = n->child (i);
      if (child->count) {
	indent (depth + 1);
	warnx << "[" << i << "]: ";
	child->dump (depth + 1);
      }
    }
  }
}

//////////////// merkle_tree_disk /////////////////

static inline str
safe_fname (str rwname)
{
  strbuf roname ("%s.ro", rwname.cstr ());
  return roname;
}

void
merkle_tree_disk::init ()
{
  if (_writer) {
    _internal = open_file (_internal_name);
    _leaf = open_file (_leaf_name);
    _index = open_file (_index_name);
    // make sure the root is created correctly:
    merkle_node *n = get_root ();
    delete n;
    write_metadata ();
  } else {
    _internal = open_file (safe_fname (_internal_name));
    _leaf = open_file (safe_fname (_leaf_name));
    _index = open_file (safe_fname (_index_name));
  }
}

void
merkle_tree_disk::close ()
{
  if (_internal) { fclose (_internal); _internal = NULL; }
  if (_leaf) { fclose (_leaf); _leaf = NULL; }
  if (_index) { fclose (_index); _index = NULL; }
}

void
merkle_tree_disk::sync (bool reopen = true)
{
  // Use lock to ensure that all copies are completed safely.
  strbuf lfname ("%s.lock", _index_name.cstr ());
  ptr<lockfile> lf = lockfile::alloc (lfname, true);

  close ();
  if (_writer) {
    // Block, copy current files to safe images.
    int r = 0;
    r = copy_file (_index_name, safe_fname (_index_name));
    r = copy_file (_leaf_name, safe_fname (_leaf_name));
    r = copy_file (_internal_name, safe_fname (_internal_name));
    // XXX Check the exit code; roll-back on failure.
  }
  if (reopen)
    init ();

  lf = NULL;
}

merkle_tree_disk::merkle_tree_disk (str path, bool writer) :
  merkle_tree (),
  _index_name (strbuf () << path << "/index.mrk"),
  _internal_name (strbuf () << path << "/internal.mrk"),
  _leaf_name (strbuf () << path << "/leaf.mrk"),
  _writer (writer)
{
  init ();
}

merkle_tree_disk::merkle_tree_disk (str index, str internal, str leaf,
				    bool writer) :
  merkle_tree (),
  _index_name (index), 
  _internal_name (internal),
  _leaf_name (leaf), 
  _writer (writer)
{
  init ();
}

merkle_tree_disk::~merkle_tree_disk ()
{
  if (_writer)
    sync (false);
  close ();
}

void
merkle_tree_disk::write_metadata ()
{
  assert (_writer);

  fseek (_index, 0, SEEK_SET);
  
  assert (_free_leafs.size () == _md.num_leaf_free);
  assert (_free_internals.size () == _md.num_internal_free);

  _md.num_leaf_free += _future_free_leafs.size ();
  _md.num_internal_free += _future_free_internals.size ();

  fwrite (&_md, sizeof(merkle_index_metadata), 1, _index);

  int nfree = _md.num_leaf_free + _md.num_internal_free;
  u_int32_t freelist[nfree];

  // write out both the unused free slots and the newly created free slots
  uint i;
  uint sofar = 0;
  for (i = 0; i < _free_leafs.size(); i++) {
    freelist[i] = ((_free_leafs[i-sofar] << 1) | 0x00000001);
  }
  sofar = i;
  for (; (i-sofar) < _free_internals.size(); i++) {
    freelist[i] = (_free_internals[i-sofar] << 1);
  }
  sofar = i;
  for (; (i-sofar) < _future_free_leafs.size(); i++) {
    freelist[i] = ((_future_free_leafs[i-sofar] << 1) | 0x00000001);
  }
  sofar = i;
  for (; (i-sofar) < _future_free_internals.size(); i++) {
    freelist[i] = (_future_free_internals[i-sofar] << 1);
  }

  fwrite (&freelist, sizeof (u_int32_t), nfree, _index);
  _future_free_leafs.clear ();
  _future_free_internals.clear ();
}

merkle_node *merkle_tree_disk::get_root ()
{
  fseek (_index, 0, SEEK_SET);

  // figure out where the root node is, given the index file
  // the first 32 bytes of the file tells us where it is
  int num_read = fread (&_md, sizeof (merkle_index_metadata), 1, _index);

  if (num_read <= 0) {
    // no root pointer yet, so we have a new tree
    _md.root = 1;
    _md.num_leaf_free = 0;
    _md.num_internal_free = 0;
    _md.next_leaf = 1;
    _md.next_internal = 0;

    // also, make a block there
    merkle_leaf_node new_root;
    bzero (&new_root, sizeof (merkle_leaf_node));
    fseek (_leaf, 0, SEEK_SET);
    fwrite (&new_root, sizeof (merkle_leaf_node), 1, _leaf);