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