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