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 #ifdef __WIN32__
  14
  15 // There's no 'ps' on win32 anyway, and Py_GetArgcArgv() isn't available.
  16 static void unpythonize_argv(void) { }
  17
  18 #else // not __WIN32__
  19
  20 // For some reason this isn't declared in Python.h
  21 extern void Py_GetArgcArgv(int *argc, char ***argv);
  22
  23 static void unpythonize_argv(void)
  24 {
  25     int argc, i;
  26     char **argv, *arge;
  27
  28     Py_GetArgcArgv(&argc, &argv);
  29
  30     for (i = 0; i < argc-1; i++)
  31     {
  32         if (argv[i] + strlen(argv[i]) + 1 != argv[i+1])
  33         {
  34             // The argv block doesn't work the way we expected; it's unsafe
  35             // to mess with it.
  36             return;
  37         }
  38     }
  39
  40     arge = argv[argc-1] + strlen(argv[argc-1]) + 1;
  41
  42     if (strstr(argv[0], "python") && argv[1] == argv[0] + strlen(argv[0]) + 1)
  43     {
  44         char *p;
  45         size_t len, diff;
  46         p = strrchr(argv[1], '/');
  47         if (p)
  48         {
  49             p++;
  50             diff = p - argv[0];
  51             len = arge - p;
  52             memmove(argv[0], p, len);
  53             memset(arge - diff, 0, diff);
  54             for (i = 0; i < argc; i++)
  55                 argv[i] = argv[i+1] ? argv[i+1]-diff : NULL;
  56         }
  57     }
  58 }
  59
  60 #endif // not __WIN32__
  61
  62
  63 static PyObject *selftest(PyObject *self, PyObject *args)
  64 {
  65     if (!PyArg_ParseTuple(args, ""))
  66         return NULL;
  67
  68     return Py_BuildValue("i", !bupsplit_selftest());
  69 }
  70
  71
  72 static PyObject *blobbits(PyObject *self, PyObject *args)
  73 {
  74     if (!PyArg_ParseTuple(args, ""))
  75         return NULL;
  76     return Py_BuildValue("i", BUP_BLOBBITS);
  77 }
  78
  79
  80 static PyObject *splitbuf(PyObject *self, PyObject *args)
  81 {
  82     unsigned char *buf = NULL;
  83     int len = 0, out = 0, bits = -1;
  84
  85     if (!PyArg_ParseTuple(args, "t#", &buf, &len))
  86         return NULL;
  87     out = bupsplit_find_ofs(buf, len, &bits);
  88     return Py_BuildValue("ii", out, bits);
  89 }
  90
  91
  92 static PyObject *bitmatch(PyObject *self, PyObject *args)
  93 {
  94     unsigned char *buf1 = NULL, *buf2 = NULL;
  95     int len1 = 0, len2 = 0;
  96     int byte, bit;
  97
  98     if (!PyArg_ParseTuple(args, "t#t#", &buf1, &len1, &buf2, &len2))
  99         return NULL;
 100
 101     bit = 0;
 102     for (byte = 0; byte < len1 && byte < len2; byte++)
 103     {
 104         int b1 = buf1[byte], b2 = buf2[byte];
 105         if (b1 != b2)
 106         {
 107             for (bit = 0; bit < 8; bit++)
 108                 if ( (b1 & (0x80 >> bit)) != (b2 & (0x80 >> bit)) )
 109                     break;
 110             break;
 111         }
 112     }
 113
 114     return Py_BuildValue("i", byte*8 + bit);
 115 }
 116
 117
 118 static PyObject *firstword(PyObject *self, PyObject *args)
 119 {
 120     unsigned char *buf = NULL;
 121     int len = 0;
 122     uint32_t v;
 123
 124     if (!PyArg_ParseTuple(args, "t#", &buf, &len))
 125         return NULL;
 126
 127     if (len < 4)
 128         return NULL;
 129
 130     v = ntohl(*(uint32_t *)buf);
 131     return PyLong_FromUnsignedLong(v);
 132 }
 133
 134
 135 typedef struct {
 136     uint32_t high;
 137     unsigned char low;
 138 } bits40_t;
 139
 140
 141 static void to_bloom_address_bitmask4(const bits40_t *buf,
 142         const int nbits, uint64_t *v, unsigned char *bitmask)
 143 {
 144     int bit;
 145     uint64_t raw, mask;
 146
 147     mask = (1<<nbits) - 1;
 148     raw = (((uint64_t)ntohl(buf->high)) << 8) | buf->low;
 149     bit = (raw >> (37-nbits)) & 0x7;
 150     *v = (raw >> (40-nbits)) & mask;
 151     *bitmask = 1 << bit;
 152 }
 153
 154 static void to_bloom_address_bitmask5(const uint32_t *buf,
 155         const int nbits, uint32_t *v, unsigned char *bitmask)
 156 {
 157     int bit;
 158     uint32_t raw, mask;
 159
 160     mask = (1<<nbits) - 1;
 161     raw = ntohl(*buf);
 162     bit = (raw >> (29-nbits)) & 0x7;
 163     *v = (raw >> (32-nbits)) & mask;
 164     *bitmask = 1 << bit;
 165 }
 166
 167
 168 #define BLOOM_SET_BIT(name, address, itype, otype) \
 169 static void name(unsigned char *bloom, const void *buf, const int nbits)\
 170 {\
 171     unsigned char bitmask;\
 172     otype v;\
 173     address((itype *)buf, nbits, &v, &bitmask);\
 174     bloom[16+v] |= bitmask;\
 175 }
 176 BLOOM_SET_BIT(bloom_set_bit4, to_bloom_address_bitmask4, bits40_t, uint64_t)
 177 BLOOM_SET_BIT(bloom_set_bit5, to_bloom_address_bitmask5, uint32_t, uint32_t)
 178
 179
 180 #define BLOOM_GET_BIT(name, address, itype, otype) \
 181 static int name(const unsigned char *bloom, const void *buf, const int nbits)\
 182 {\
 183     unsigned char bitmask;\
 184     otype v;\
 185     address((itype *)buf, nbits, &v, &bitmask);\
 186     return bloom[16+v] & bitmask;\
 187 }
 188 BLOOM_GET_BIT(bloom_get_bit4, to_bloom_address_bitmask4, bits40_t, uint64_t)
 189 BLOOM_GET_BIT(bloom_get_bit5, to_bloom_address_bitmask5, uint32_t, uint32_t)
 190
 191
 192 static PyObject *bloom_add(PyObject *self, PyObject *args)
 193 {
 194     unsigned char *sha = NULL, *bloom = NULL;
 195     unsigned char *end;
 196     int len = 0, blen = 0, nbits = 0, k = 0;
 197
 198     if (!PyArg_ParseTuple(args, "w#s#ii", &bloom, &blen, &sha, &len, &nbits, &k))
 199         return NULL;
 200
 201     if (blen < 16+(1<<nbits) || len % 20 != 0)
 202         return NULL;
 203
 204     if (k == 5)
 205     {
 206         if (nbits > 29)
 207             return NULL;
 208         for (end = sha + len; sha < end; sha += 20/k)
 209             bloom_set_bit5(bloom, sha, nbits);
 210     }
 211     else if (k == 4)
 212     {
 213         if (nbits > 37)
 214             return NULL;
 215         for (end = sha + len; sha < end; sha += 20/k)
 216             bloom_set_bit4(bloom, sha, nbits);
 217     }
 218     else
 219         return NULL;
 220
 221
 222     return Py_BuildValue("i", len/20);
 223 }
 224
 225 static PyObject *bloom_contains(PyObject *self, PyObject *args)
 226 {
 227     unsigned char *sha = NULL, *bloom = NULL;
 228     int len = 0, blen = 0, nbits = 0, k = 0;
 229     unsigned char *end;
 230     int steps;
 231
 232     if (!PyArg_ParseTuple(args, "t#s#ii", &bloom, &blen, &sha, &len, &nbits, &k))
 233         return NULL;
 234
 235     if (len != 20)
 236         return NULL;
 237
 238     if (k == 5)
 239     {
 240         if (nbits > 29)
 241             return NULL;
 242         for (steps = 1, end = sha + 20; sha < end; sha += 20/k, steps++)
 243             if (!bloom_get_bit5(bloom, sha, nbits))
 244                 return Py_BuildValue("Oi", Py_None, steps);
 245     }
 246     else if (k == 4)
 247     {
 248         if (nbits > 37)
 249             return NULL;
 250         for (steps = 1, end = sha + 20; sha < end; sha += 20/k, steps++)
 251             if (!bloom_get_bit4(bloom, sha, nbits))
 252                 return Py_BuildValue("Oi", Py_None, steps);
 253     }
 254     else
 255         return NULL;
 256
 257     return Py_BuildValue("Oi", Py_True, k);
 258 }
 259
 260
 261 static uint32_t _extract_bits(unsigned char *buf, int nbits)
 262 {
 263     uint32_t v, mask;
 264
 265     mask = (1<<nbits) - 1;
 266     v = ntohl(*(uint32_t *)buf);
 267     v = (v >> (32-nbits)) & mask;
 268     return v;
 269 }
 270 static PyObject *extract_bits(PyObject *self, PyObject *args)
 271 {
 272     unsigned char *buf = NULL;
 273     int len = 0, nbits = 0;
 274
 275     if (!PyArg_ParseTuple(args, "t#i", &buf, &len, &nbits))
 276         return NULL;
 277
 278     if (len < 4)
 279         return NULL;
 280
 281     return PyLong_FromUnsignedLong(_extract_bits(buf, nbits));
 282 }
 283
 284
 285 struct sha {
 286     unsigned char bytes[20];
 287 };
 288 struct idx {
 289     unsigned char *map;
 290     struct sha *cur;
 291     struct sha *end;
 292     uint32_t *cur_name;
 293     long bytes;
 294     int name_base;
 295 };
 296
 297
 298 static int _cmp_sha(const struct sha *sha1, const struct sha *sha2)
 299 {
 300     int i;
 301     for (i = 0; i < 20; i++)
 302         if (sha1->bytes[i] != sha2->bytes[i])
 303             return sha1->bytes[i] - sha2->bytes[i];
 304     return 0;
 305 }
 306
 307
 308 static void _fix_idx_order(struct idx **idxs, int *last_i)
 309 {
 310     struct idx *idx;
 311     int low, mid, high, c = 0;
 312
 313     idx = idxs[*last_i];
 314     if (idxs[*last_i]->cur >= idxs[*last_i]->end)
 315     {
 316         idxs[*last_i] = NULL;
 317         PyMem_Free(idx);
 318         --*last_i;
 319         return;
 320     }
 321     if (*last_i == 0)
 322         return;
 323
 324     low = *last_i-1;
 325     mid = *last_i;
 326     high = 0;
 327     while (low >= high)
 328     {
 329         mid = (low + high) / 2;
 330         c = _cmp_sha(idx->cur, idxs[mid]->cur);
 331         if (c < 0)
 332             high = mid + 1;
 333         else if (c > 0)
 334             low = mid - 1;
 335         else
 336             break;
 337     }
 338     if (c < 0)
 339         ++mid;
 340     if (mid == *last_i)
 341         return;
 342     memmove(&idxs[mid+1], &idxs[mid], (*last_i-mid)*sizeof(struct idx *));
 343     idxs[mid] = idx;
 344 }
 345
 346
 347 static uint32_t _get_idx_i(struct idx *idx)
 348 {
 349     if (idx->cur_name == NULL)
 350         return idx->name_base;
 351     return ntohl(*idx->cur_name) + idx->name_base;
 352 }
 353
 354
 355 static PyObject *merge_into(PyObject *self, PyObject *args)
 356 {
 357     PyObject *ilist = NULL;
 358     unsigned char *fmap = NULL;
 359     struct sha *sha_ptr, *last = NULL;
 360     uint32_t *table_ptr, *name_ptr;
 361     struct idx **idxs = NULL;
 362     int flen = 0, bits = 0, i;
 363     uint32_t total, count, prefix;
 364     int num_i;
 365     int last_i;
 366
 367     if (!PyArg_ParseTuple(args, "w#iIO", &fmap, &flen, &bits, &total, &ilist))
 368         return NULL;
 369
 370     num_i = PyList_Size(ilist);
 371     idxs = (struct idx **)PyMem_Malloc(num_i * sizeof(struct idx *));
 372
 373     for (i = 0; i < num_i; i++)
 374     {
 375         long len, sha_ofs, name_map_ofs;
 376         idxs[i] = (struct idx *)PyMem_Malloc(sizeof(struct idx));
 377         PyObject *itup = PyList_GetItem(ilist, i);
 378         if (!PyArg_ParseTuple(itup, "t#llli", &idxs[i]->map, &idxs[i]->bytes,
 379                     &len, &sha_ofs, &name_map_ofs, &idxs[i]->name_base))
 380             return NULL;
 381         idxs[i]->cur = (struct sha *)&idxs[i]->map[sha_ofs];
 382         idxs[i]->end = &idxs[i]->cur[len];
 383         if (name_map_ofs)
 384             idxs[i]->cur_name = (uint32_t *)&idxs[i]->map[name_map_ofs];
 385         else
 386             idxs[i]->cur_name = NULL;
 387     }
 388     table_ptr = (uint32_t *)&fmap[12];
 389     sha_ptr = (struct sha *)&table_ptr[1<<bits];
 390     name_ptr = (uint32_t *)&sha_ptr[total];
 391
 392     last_i = num_i-1;
 393     count = 0;
 394     prefix = 0;
 395     while (last_i >= 0)
 396     {
 397         struct idx *idx;
 398         uint32_t new_prefix;
 399         if (count % 102424 == 0 && istty)
 400             fprintf(stderr, "midx: writing %.2f%% (%d/%d)\r",
 401                     count*100.0/total, count, total);
 402         idx = idxs[last_i];
 403         new_prefix = _extract_bits((unsigned char *)idx->cur, bits);
 404         while (prefix < new_prefix)
 405             table_ptr[prefix++] = htonl(count);
 406         if (last == NULL || _cmp_sha(last, idx->cur) != 0)
 407         {
 408             memcpy(sha_ptr++, idx->cur, 20);
 409             *name_ptr++ = htonl(_get_idx_i(idx));
 410             last = idx->cur;
 411         }
 412         ++idx->cur;
 413         if (idx->cur_name != NULL)
 414             ++idx->cur_name;
 415         _fix_idx_order(idxs, &last_i);
 416         ++count;
 417     }
 418     table_ptr[prefix] = htonl(count);
 419
 420     PyMem_Free(idxs);
 421     return PyLong_FromUnsignedLong(count);
 422 }
 423
 424 // This function should techniclly be macro'd out if it's going to be used
 425 // more than ocasionally.  As of this writing, it'll actually never be called
 426 // in real world bup scenarios (because our packs are < MAX_INT bytes).
 427 static uint64_t htonll(uint64_t value)
 428 {
 429     static const int endian_test = 42;
 430
 431     if (*(char *)&endian_test == endian_test) // LSB-MSB
 432         return ((uint64_t)htonl(value & 0xFFFFFFFF) << 32) | htonl(value >> 32);
 433     return value; // already in network byte order MSB-LSB
 434 }
 435
 436 static PyObject *write_idx(PyObject *self, PyObject *args)
 437 {
 438     PyObject *pf = NULL, *idx = NULL;
 439     PyObject *part;
 440     FILE *f;
 441     unsigned char *fmap = NULL;
 442     int flen = 0;
 443     uint32_t total = 0;
 444     uint32_t count;
 445     int i, j, ofs64_count;
 446     uint32_t *fan_ptr, *crc_ptr, *ofs_ptr;
 447     struct sha *sha_ptr;
 448
 449     if (!PyArg_ParseTuple(args, "Ow#OI", &pf, &fmap, &flen, &idx, &total))
 450         return NULL;
 451
 452     fan_ptr = (uint32_t *)&fmap[8];
 453     sha_ptr = (struct sha *)&fan_ptr[256];
 454     crc_ptr = (uint32_t *)&sha_ptr[total];
 455     ofs_ptr = (uint32_t *)&crc_ptr[total];
 456     f = PyFile_AsFile(pf);
 457
 458     count = 0;
 459     ofs64_count = 0;
 460     for (i = 0; i < 256; ++i)
 461     {
 462         int plen;
 463         part = PyList_GET_ITEM(idx, i);
 464         PyList_Sort(part);
 465         plen = PyList_GET_SIZE(part);
 466         count += plen;
 467         *fan_ptr++ = htonl(count);
 468         for (j = 0; j < plen; ++j)
 469         {
 470             unsigned char *sha = NULL;
 471             int sha_len = 0;
 472             uint32_t crc = 0;
 473             uint64_t ofs = 0;
 474             if (!PyArg_ParseTuple(PyList_GET_ITEM(part, j), "t#IK",
 475                                   &sha, &sha_len, &crc, &ofs))
 476                 return NULL;
 477             if (sha_len != 20)
 478                 return NULL;
 479             memcpy(sha_ptr++, sha, 20);
 480             *crc_ptr++ = htonl(crc);
 481             if (ofs > 0x7fffffff)
 482             {
 483                 uint64_t nofs = htonll(ofs);
 484                 if (fwrite(&nofs, 8, 1, f) != 1)
 485                     return PyErr_SetFromErrno(PyExc_OSError);
 486                 ofs = 0x80000000 | ofs64_count++;
 487             }
 488             *ofs_ptr++ = htonl((uint32_t)ofs);
 489         }
 490     }
 491     return PyLong_FromUnsignedLong(count);
 492 }
 493
 494
 495 // I would have made this a lower-level function that just fills in a buffer
 496 // with random values, and then written those values from python.  But that's
 497 // about 20% slower in my tests, and since we typically generate random
 498 // numbers for benchmarking other parts of bup, any slowness in generating
 499 // random bytes will make our benchmarks inaccurate.  Plus nobody wants
 500 // pseudorandom bytes much except for this anyway.
 501 static PyObject *write_random(PyObject *self, PyObject *args)
 502 {
 503     uint32_t buf[1024/4];
 504     int fd = -1, seed = 0, verbose = 0;
 505     ssize_t ret;
 506     long long len = 0, kbytes = 0, written = 0;
 507
 508     if (!PyArg_ParseTuple(args, "iLii", &fd, &len, &seed, &verbose))
 509         return NULL;
 510
 511     srandom(seed);
 512
 513     for (kbytes = 0; kbytes < len/1024; kbytes++)
 514     {
 515         unsigned i;
 516         for (i = 0; i < sizeof(buf)/sizeof(buf[0]); i++)
 517             buf[i] = random();
 518         ret = write(fd, buf, sizeof(buf));
 519         if (ret < 0)
 520             ret = 0;
 521         written += ret;
 522         if (ret < (int)sizeof(buf))
 523             break;
 524         if (verbose && kbytes/1024 > 0 && !(kbytes%1024))
 525             fprintf(stderr, "Random: %lld Mbytes\r", kbytes/1024);
 526     }
 527
 528     // handle non-multiples of 1024
 529     if (len % 1024)
 530     {
 531         unsigned i;
 532         for (i = 0; i < sizeof(buf)/sizeof(buf[0]); i++)
 533             buf[i] = random();
 534         ret = write(fd, buf, len % 1024);
 535         if (ret < 0)
 536             ret = 0;
 537         written += ret;
 538     }
 539
 540     if (kbytes/1024 > 0)
 541         fprintf(stderr, "Random: %lld Mbytes, done.\n", kbytes/1024);
 542     return Py_BuildValue("L", written);
 543 }
 544
 545
 546 static PyObject *random_sha(PyObject *self, PyObject *args)
 547 {
 548     static int seeded = 0;
 549     uint32_t shabuf[20/4];
 550     int i;
 551
 552     if (!seeded)
 553     {
 554         assert(sizeof(shabuf) == 20);
 555         srandom(time(NULL));
 556         seeded = 1;
 557     }
 558
 559     if (!PyArg_ParseTuple(args, ""))
 560         return NULL;
 561
 562     memset(shabuf, 0, sizeof(shabuf));
 563     for (i=0; i < 20/4; i++)
 564         shabuf[i] = random();
 565     return Py_BuildValue("s#", shabuf, 20);
 566 }
 567
 568
 569 static PyObject *open_noatime(PyObject *self, PyObject *args)
 570 {
 571     char *filename = NULL;
 572     int attrs, attrs_noatime, fd;
 573     if (!PyArg_ParseTuple(args, "s", &filename))
 574         return NULL;
 575     attrs = O_RDONLY;
 576 #ifdef O_NOFOLLOW
 577     attrs |= O_NOFOLLOW;
 578 #endif
 579 #ifdef O_LARGEFILE
 580     attrs |= O_LARGEFILE;
 581 #endif
 582     attrs_noatime = attrs;
 583 #ifdef O_NOATIME
 584     attrs_noatime |= O_NOATIME;
 585 #endif
 586     fd = open(filename, attrs_noatime);
 587     if (fd < 0 && errno == EPERM)
 588     {
 589         // older Linux kernels would return EPERM if you used O_NOATIME
 590         // and weren't the file's owner.  This pointless restriction was
 591         // relaxed eventually, but we have to handle it anyway.
 592         // (VERY old kernels didn't recognized O_NOATIME, but they would
 593         // just harmlessly ignore it, so this branch won't trigger)
 594         fd = open(filename, attrs);
 595     }
 596     if (fd < 0)
 597         return PyErr_SetFromErrnoWithFilename(PyExc_IOError, filename);
 598     return Py_BuildValue("i", fd);
 599 }
 600
 601
 602 static PyObject *fadvise_done(PyObject *self, PyObject *args)
 603 {
 604     int fd = -1;
 605     long long ofs = 0;
 606     if (!PyArg_ParseTuple(args, "iL", &fd, &ofs))
 607         return NULL;
 608 #ifdef POSIX_FADV_DONTNEED
 609     posix_fadvise(fd, 0, ofs, POSIX_FADV_DONTNEED);
 610 #endif
 611     return Py_BuildValue("");
 612 }
 613
 614
 615 static PyMethodDef faster_methods[] = {
 616     { "selftest", selftest, METH_VARARGS,
 617         "Check that the rolling checksum rolls correctly (for unit tests)." },
 618     { "blobbits", blobbits, METH_VARARGS,
 619         "Return the number of bits in the rolling checksum." },
 620     { "splitbuf", splitbuf, METH_VARARGS,
 621         "Split a list of strings based on a rolling checksum." },
 622     { "bitmatch", bitmatch, METH_VARARGS,
 623         "Count the number of matching prefix bits between two strings." },
 624     { "firstword", firstword, METH_VARARGS,
 625         "Return an int corresponding to the first 32 bits of buf." },
 626     { "bloom_contains", bloom_contains, METH_VARARGS,
 627         "Check if a bloom filter of 2^nbits bytes contains an object" },
 628     { "bloom_add", bloom_add, METH_VARARGS,
 629         "Add an object to a bloom filter of 2^nbits bytes" },
 630     { "extract_bits", extract_bits, METH_VARARGS,
 631         "Take the first 'nbits' bits from 'buf' and return them as an int." },
 632     { "merge_into", merge_into, METH_VARARGS,
 633         "Merges a bunch of idx and midx files into a single midx." },
 634     { "write_idx", write_idx, METH_VARARGS,
 635         "Write a PackIdxV2 file from an idx list of lists of tuples" },
 636     { "write_random", write_random, METH_VARARGS,
 637         "Write random bytes to the given file descriptor" },
 638     { "random_sha", random_sha, METH_VARARGS,
 639         "Return a random 20-byte string" },
 640     { "open_noatime", open_noatime, METH_VARARGS,
 641         "open() the given filename for read with O_NOATIME if possible" },
 642     { "fadvise_done", fadvise_done, METH_VARARGS,
 643         "Inform the kernel that we're finished with earlier parts of a file" },
 644     { NULL, NULL, 0, NULL },  // sentinel
 645 };
 646
 647 PyMODINIT_FUNC init_helpers(void)
 648 {
 649     Py_InitModule("_helpers", faster_methods);
 650     istty = isatty(2) || getenv("BUP_FORCE_TTY");
 651     unpythonize_argv();
 652 }