lib/bup/_helpers.c

   1 #include "bupsplit.h"
   2 #include <Python.h>
   3 #include <assert.h>
   4 #include <stdint.h>
   5 #include <fcntl.h>
   6 #include <arpa/inet.h>
   7 #include <unistd.h>
   8 #include <stdlib.h>
   9 #include <stdio.h>
  10
  11 static int istty = 0;
  12
  13 static PyObject *selftest(PyObject *self, PyObject *args)
  14 {
  15     if (!PyArg_ParseTuple(args, ""))
  16         return NULL;
  17
  18     return Py_BuildValue("i", !bupsplit_selftest());
  19 }
  20
  21
  22 static PyObject *blobbits(PyObject *self, PyObject *args)
  23 {
  24     if (!PyArg_ParseTuple(args, ""))
  25         return NULL;
  26     return Py_BuildValue("i", BUP_BLOBBITS);
  27 }
  28
  29
  30 static PyObject *splitbuf(PyObject *self, PyObject *args)
  31 {
  32     unsigned char *buf = NULL;
  33     int len = 0, out = 0, bits = -1;
  34
  35     if (!PyArg_ParseTuple(args, "t#", &buf, &len))
  36         return NULL;
  37     out = bupsplit_find_ofs(buf, len, &bits);
  38     return Py_BuildValue("ii", out, bits);
  39 }
  40
  41
  42 static PyObject *bitmatch(PyObject *self, PyObject *args)
  43 {
  44     unsigned char *buf1 = NULL, *buf2 = NULL;
  45     int len1 = 0, len2 = 0;
  46     int byte, bit;
  47
  48     if (!PyArg_ParseTuple(args, "t#t#", &buf1, &len1, &buf2, &len2))
  49         return NULL;
  50
  51     bit = 0;
  52     for (byte = 0; byte < len1 && byte < len2; byte++)
  53     {
  54         int b1 = buf1[byte], b2 = buf2[byte];
  55         if (b1 != b2)
  56         {
  57             for (bit = 0; bit < 8; bit++)
  58                 if ( (b1 & (0x80 >> bit)) != (b2 & (0x80 >> bit)) )
  59                     break;
  60             break;
  61         }
  62     }
  63
  64     return Py_BuildValue("i", byte*8 + bit);
  65 }
  66
  67
  68 static PyObject *firstword(PyObject *self, PyObject *args)
  69 {
  70     unsigned char *buf = NULL;
  71     int len = 0;
  72     uint32_t v;
  73
  74     if (!PyArg_ParseTuple(args, "t#", &buf, &len))
  75         return NULL;
  76
  77     if (len < 4)
  78         return NULL;
  79
  80     v = ntohl(*(uint32_t *)buf);
  81     return PyLong_FromUnsignedLong(v);
  82 }
  83
  84
  85 typedef struct {
  86     uint32_t high;
  87     unsigned char low;
  88 } bits40_t;
  89
  90
  91 static void to_bloom_address_bitmask4(const bits40_t *buf,
  92         const int nbits, uint64_t *v, unsigned char *bitmask)
  93 {
  94     int bit;
  95     uint64_t raw, mask;
  96
  97     mask = (1<<nbits) - 1;
  98     raw = (((uint64_t)ntohl(buf->high)) << 8) | buf->low;
  99     bit = (raw >> (37-nbits)) & 0x7;
 100     *v = (raw >> (40-nbits)) & mask;
 101     *bitmask = 1 << bit;
 102 }
 103
 104 static void to_bloom_address_bitmask5(const uint32_t *buf,
 105         const int nbits, uint32_t *v, unsigned char *bitmask)
 106 {
 107     int bit;
 108     uint32_t raw, mask;
 109
 110     mask = (1<<nbits) - 1;
 111     raw = ntohl(*buf);
 112     bit = (raw >> (29-nbits)) & 0x7;
 113     *v = (raw >> (32-nbits)) & mask;
 114     *bitmask = 1 << bit;
 115 }
 116
 117
 118 #define BLOOM_SET_BIT(name, address, itype, otype) \
 119 static void name(unsigned char *bloom, const void *buf, const int nbits)\
 120 {\
 121     unsigned char bitmask;\
 122     otype v;\
 123     address((itype *)buf, nbits, &v, &bitmask);\
 124     bloom[16+v] |= bitmask;\
 125 }
 126 BLOOM_SET_BIT(bloom_set_bit4, to_bloom_address_bitmask4, bits40_t, uint64_t)
 127 BLOOM_SET_BIT(bloom_set_bit5, to_bloom_address_bitmask5, uint32_t, uint32_t)
 128
 129
 130 #define BLOOM_GET_BIT(name, address, itype, otype) \
 131 static int name(const unsigned char *bloom, const void *buf, const int nbits)\
 132 {\
 133     unsigned char bitmask;\
 134     otype v;\
 135     address((itype *)buf, nbits, &v, &bitmask);\
 136     return bloom[16+v] & bitmask;\
 137 }
 138 BLOOM_GET_BIT(bloom_get_bit4, to_bloom_address_bitmask4, bits40_t, uint64_t)
 139 BLOOM_GET_BIT(bloom_get_bit5, to_bloom_address_bitmask5, uint32_t, uint32_t)
 140
 141
 142 static PyObject *bloom_add(PyObject *self, PyObject *args)
 143 {
 144     unsigned char *sha = NULL, *bloom = NULL;
 145     unsigned char *end;
 146     int len = 0, blen = 0, nbits = 0, k = 0;
 147
 148     if (!PyArg_ParseTuple(args, "w#s#ii", &bloom, &blen, &sha, &len, &nbits, &k))
 149         return NULL;
 150
 151     if (blen < 16+(1<<nbits) || len % 20 != 0)
 152         return NULL;
 153
 154     if (k == 5)
 155     {
 156         if (nbits > 29)
 157             return NULL;
 158         for (end = sha + len; sha < end; sha += 20/k)
 159             bloom_set_bit5(bloom, sha, nbits);
 160     }
 161     else if (k == 4)
 162     {
 163         if (nbits > 37)
 164             return NULL;
 165         for (end = sha + len; sha < end; sha += 20/k)
 166             bloom_set_bit4(bloom, sha, nbits);
 167     }
 168     else
 169         return NULL;
 170
 171
 172     return Py_BuildValue("i", len/20);
 173 }
 174
 175 static PyObject *bloom_contains(PyObject *self, PyObject *args)
 176 {
 177     unsigned char *sha = NULL, *bloom = NULL;
 178     int len = 0, blen = 0, nbits = 0, k = 0;
 179     unsigned char *end;
 180     int steps;
 181
 182     if (!PyArg_ParseTuple(args, "t#s#ii", &bloom, &blen, &sha, &len, &nbits, &k))
 183         return NULL;
 184
 185     if (len != 20)
 186         return NULL;
 187
 188     if (k == 5)
 189     {
 190         if (nbits > 29)
 191             return NULL;
 192         for (steps = 1, end = sha + 20; sha < end; sha += 20/k, steps++)
 193             if (!bloom_get_bit5(bloom, sha, nbits))
 194                 return Py_BuildValue("Oi", Py_None, steps);
 195     }
 196     else if (k == 4)
 197     {
 198         if (nbits > 37)
 199             return NULL;
 200         for (steps = 1, end = sha + 20; sha < end; sha += 20/k, steps++)
 201             if (!bloom_get_bit4(bloom, sha, nbits))
 202                 return Py_BuildValue("Oi", Py_None, steps);
 203     }
 204     else
 205         return NULL;
 206
 207     return Py_BuildValue("Oi", Py_True, k);
 208 }
 209
 210
 211 static uint32_t _extract_bits(unsigned char *buf, int nbits)
 212 {
 213     uint32_t v, mask;
 214
 215     mask = (1<<nbits) - 1;
 216     v = ntohl(*(uint32_t *)buf);
 217     v = (v >> (32-nbits)) & mask;
 218     return v;
 219 }
 220 static PyObject *extract_bits(PyObject *self, PyObject *args)
 221 {
 222     unsigned char *buf = NULL;
 223     int len = 0, nbits = 0;
 224
 225     if (!PyArg_ParseTuple(args, "t#i", &buf, &len, &nbits))
 226         return NULL;
 227
 228     if (len < 4)
 229         return NULL;
 230
 231     return PyLong_FromUnsignedLong(_extract_bits(buf, nbits));
 232 }
 233
 234
 235 struct sha {
 236     unsigned char bytes[20];
 237 };
 238 struct idx {
 239     unsigned char *map;
 240     struct sha *cur;
 241     struct sha *end;
 242     uint32_t *cur_name;
 243     long bytes;
 244     int name_base;
 245 };
 246
 247
 248 static int _cmp_sha(const struct sha *sha1, const struct sha *sha2)
 249 {
 250     int i;
 251     for (i = 0; i < 20; i++)
 252         if (sha1->bytes[i] != sha2->bytes[i])
 253             return sha1->bytes[i] - sha2->bytes[i];
 254     return 0;
 255 }
 256
 257
 258 static void _fix_idx_order(struct idx **idxs, int *last_i)
 259 {
 260     struct idx *idx;
 261     int low, mid, high, c = 0;
 262
 263     idx = idxs[*last_i];
 264     if (idxs[*last_i]->cur >= idxs[*last_i]->end)
 265     {
 266         idxs[*last_i] = NULL;
 267         PyMem_Free(idx);
 268         --*last_i;
 269         return;
 270     }
 271     if (*last_i == 0)
 272         return;
 273
 274     low = *last_i-1;
 275     mid = *last_i;
 276     high = 0;
 277     while (low >= high)
 278     {
 279         mid = (low + high) / 2;
 280         c = _cmp_sha(idx->cur, idxs[mid]->cur);
 281         if (c < 0)
 282             high = mid + 1;
 283         else if (c > 0)
 284             low = mid - 1;
 285         else
 286             break;
 287     }
 288     if (c < 0)
 289         ++mid;
 290     if (mid == *last_i)
 291         return;
 292     memmove(&idxs[mid+1], &idxs[mid], (*last_i-mid)*sizeof(struct idx *));
 293     idxs[mid] = idx;
 294 }
 295
 296
 297 static uint32_t _get_idx_i(struct idx *idx)
 298 {
 299     if (idx->cur_name == NULL)
 300         return idx->name_base;
 301     return ntohl(*idx->cur_name) + idx->name_base;
 302 }
 303
 304
 305 static PyObject *merge_into(PyObject *self, PyObject *args)
 306 {
 307     PyObject *ilist = NULL;
 308     unsigned char *fmap = NULL;
 309     struct sha *sha_ptr, *last = NULL;
 310     uint32_t *table_ptr, *name_ptr;
 311     struct idx **idxs = NULL;
 312     int flen = 0, bits = 0, i;
 313     uint32_t total, count, prefix;
 314     int num_i;
 315     int last_i;
 316
 317     if (!PyArg_ParseTuple(args, "w#iIO", &fmap, &flen, &bits, &total, &ilist))
 318         return NULL;
 319
 320     num_i = PyList_Size(ilist);
 321     idxs = (struct idx **)PyMem_Malloc(num_i * sizeof(struct idx *));
 322
 323     for (i = 0; i < num_i; i++)
 324     {
 325         long len, sha_ofs, name_map_ofs;
 326         idxs[i] = (struct idx *)PyMem_Malloc(sizeof(struct idx));
 327         PyObject *itup = PyList_GetItem(ilist, i);
 328         if (!PyArg_ParseTuple(itup, "t#llli", &idxs[i]->map, &idxs[i]->bytes,
 329                     &len, &sha_ofs, &name_map_ofs, &idxs[i]->name_base))
 330             return NULL;
 331         idxs[i]->cur = (struct sha *)&idxs[i]->map[sha_ofs];
 332         idxs[i]->end = &idxs[i]->cur[len];
 333         if (name_map_ofs)
 334             idxs[i]->cur_name = (uint32_t *)&idxs[i]->map[name_map_ofs];
 335         else
 336             idxs[i]->cur_name = NULL;
 337     }
 338     table_ptr = (uint32_t *)&fmap[12];
 339     sha_ptr = (struct sha *)&table_ptr[1<<bits];
 340     name_ptr = (uint32_t *)&sha_ptr[total];
 341
 342     last_i = num_i-1;
 343     count = 0;
 344     prefix = 0;
 345     while (last_i >= 0)
 346     {
 347         struct idx *idx;
 348         uint32_t new_prefix;
 349         if (count % 102424 == 0 && istty)
 350             fprintf(stderr, "midx: writing %.2f%% (%d/%d)\r",
 351                     count*100.0/total, count, total);
 352         idx = idxs[last_i];
 353         new_prefix = _extract_bits((unsigned char *)idx->cur, bits);
 354         while (prefix < new_prefix)
 355             table_ptr[prefix++] = htonl(count);
 356         if (last == NULL || _cmp_sha(last, idx->cur) != 0)
 357         {
 358             memcpy(sha_ptr++, idx->cur, 20);
 359             *name_ptr++ = htonl(_get_idx_i(idx));
 360             last = idx->cur;
 361         }
 362         ++idx->cur;
 363         if (idx->cur_name != NULL)
 364             ++idx->cur_name;
 365         _fix_idx_order(idxs, &last_i);
 366         ++count;
 367     }
 368     table_ptr[prefix] = htonl(count);
 369
 370     PyMem_Free(idxs);
 371     return PyLong_FromUnsignedLong(count);
 372 }
 373
 374
 375 // I would have made this a lower-level function that just fills in a buffer
 376 // with random values, and then written those values from python.  But that's
 377 // about 20% slower in my tests, and since we typically generate random
 378 // numbers for benchmarking other parts of bup, any slowness in generating
 379 // random bytes will make our benchmarks inaccurate.  Plus nobody wants
 380 // pseudorandom bytes much except for this anyway.
 381 static PyObject *write_random(PyObject *self, PyObject *args)
 382 {
 383     uint32_t buf[1024/4];
 384     int fd = -1, seed = 0, verbose = 0;
 385     ssize_t ret;
 386     long long len = 0, kbytes = 0, written = 0;
 387
 388     if (!PyArg_ParseTuple(args, "iLii", &fd, &len, &seed, &verbose))
 389         return NULL;
 390
 391     srandom(seed);
 392
 393     for (kbytes = 0; kbytes < len/1024; kbytes++)
 394     {
 395         unsigned i;
 396         for (i = 0; i < sizeof(buf)/sizeof(buf[0]); i++)
 397             buf[i] = random();
 398         ret = write(fd, buf, sizeof(buf));
 399         if (ret < 0)
 400             ret = 0;
 401         written += ret;
 402         if (ret < (int)sizeof(buf))
 403             break;
 404         if (verbose && kbytes/1024 > 0 && !(kbytes%1024))
 405             fprintf(stderr, "Random: %lld Mbytes\r", kbytes/1024);
 406     }
 407
 408     // handle non-multiples of 1024
 409     if (len % 1024)
 410     {
 411         unsigned i;
 412         for (i = 0; i < sizeof(buf)/sizeof(buf[0]); i++)
 413             buf[i] = random();
 414         ret = write(fd, buf, len % 1024);
 415         if (ret < 0)
 416             ret = 0;
 417         written += ret;
 418     }
 419
 420     if (kbytes/1024 > 0)
 421         fprintf(stderr, "Random: %lld Mbytes, done.\n", kbytes/1024);
 422     return Py_BuildValue("L", written);
 423 }
 424
 425
 426 static PyObject *random_sha(PyObject *self, PyObject *args)
 427 {
 428     static int seeded = 0;
 429     uint32_t shabuf[20/4];
 430     int i;
 431
 432     if (!seeded)
 433     {
 434         assert(sizeof(shabuf) == 20);
 435         srandom(time(NULL));
 436         seeded = 1;
 437     }
 438
 439     if (!PyArg_ParseTuple(args, ""))
 440         return NULL;
 441
 442     memset(shabuf, 0, sizeof(shabuf));
 443     for (i=0; i < 20/4; i++)
 444         shabuf[i] = random();
 445     return Py_BuildValue("s#", shabuf, 20);
 446 }
 447
 448
 449 static PyObject *open_noatime(PyObject *self, PyObject *args)
 450 {
 451     char *filename = NULL;
 452     int attrs, attrs_noatime, fd;
 453     if (!PyArg_ParseTuple(args, "s", &filename))
 454         return NULL;
 455     attrs = O_RDONLY;
 456 #ifdef O_NOFOLLOW
 457     attrs |= O_NOFOLLOW;
 458 #endif
 459 #ifdef O_LARGEFILE
 460     attrs |= O_LARGEFILE;
 461 #endif
 462     attrs_noatime = attrs;
 463 #ifdef O_NOATIME
 464     attrs_noatime |= O_NOATIME;
 465 #endif
 466     fd = open(filename, attrs_noatime);
 467     if (fd < 0 && errno == EPERM)
 468     {
 469         // older Linux kernels would return EPERM if you used O_NOATIME
 470         // and weren't the file's owner.  This pointless restriction was
 471         // relaxed eventually, but we have to handle it anyway.
 472         // (VERY old kernels didn't recognized O_NOATIME, but they would
 473         // just harmlessly ignore it, so this branch won't trigger)
 474         fd = open(filename, attrs);
 475     }
 476     if (fd < 0)
 477         return PyErr_SetFromErrnoWithFilename(PyExc_IOError, filename);
 478     return Py_BuildValue("i", fd);
 479 }
 480
 481
 482 static PyObject *fadvise_done(PyObject *self, PyObject *args)
 483 {
 484     int fd = -1;
 485     long long ofs = 0;
 486     if (!PyArg_ParseTuple(args, "iL", &fd, &ofs))
 487         return NULL;
 488 #ifdef POSIX_FADV_DONTNEED
 489     posix_fadvise(fd, 0, ofs, POSIX_FADV_DONTNEED);
 490 #endif
 491     return Py_BuildValue("");
 492 }
 493
 494
 495 static PyMethodDef faster_methods[] = {
 496     { "selftest", selftest, METH_VARARGS,
 497         "Check that the rolling checksum rolls correctly (for unit tests)." },
 498     { "blobbits", blobbits, METH_VARARGS,
 499         "Return the number of bits in the rolling checksum." },
 500     { "splitbuf", splitbuf, METH_VARARGS,
 501         "Split a list of strings based on a rolling checksum." },
 502     { "bitmatch", bitmatch, METH_VARARGS,
 503         "Count the number of matching prefix bits between two strings." },
 504     { "firstword", firstword, METH_VARARGS,
 505         "Return an int corresponding to the first 32 bits of buf." },
 506     { "bloom_contains", bloom_contains, METH_VARARGS,
 507         "Check if a bloom filter of 2^nbits bytes contains an object" },
 508     { "bloom_add", bloom_add, METH_VARARGS,
 509         "Add an object to a bloom filter of 2^nbits bytes" },
 510     { "extract_bits", extract_bits, METH_VARARGS,
 511         "Take the first 'nbits' bits from 'buf' and return them as an int." },
 512     { "merge_into", merge_into, METH_VARARGS,
 513         "Merges a bunch of idx and midx files into a single midx." },
 514     { "write_random", write_random, METH_VARARGS,
 515         "Write random bytes to the given file descriptor" },
 516     { "random_sha", random_sha, METH_VARARGS,
 517         "Return a random 20-byte string" },
 518     { "open_noatime", open_noatime, METH_VARARGS,
 519         "open() the given filename for read with O_NOATIME if possible" },
 520     { "fadvise_done", fadvise_done, METH_VARARGS,
 521         "Inform the kernel that we're finished with earlier parts of a file" },
 522     { NULL, NULL, 0, NULL },  // sentinel
 523 };
 524
 525 PyMODINIT_FUNC init_helpers(void)
 526 {
 527     Py_InitModule("_helpers", faster_methods);
 528     istty = isatty(2) || getenv("BUP_FORCE_TTY");
 529 }