ucw/ddigger.c

   1 #include <ucw/lib.h>
   2 #include <ucw/io.h>
   3 #include <ucw/opt.h>
   4 #include <stdint.h>
   5 #include <stdio.h>
   6 #include <fcntl.h>
   7 #include <unistd.h>
   8 #include <sys/ioctl.h>
   9 #include <linux/fs.h>
  10
  11 static char *dev_name;
  12 static char *out_name;
  13 static char *stat_name;
  14 static uint block_size = 1048576;
  15 static uint error_skip;
  16 static u64 start_pos;
  17 static u64 end_pos = ~(u64)0;
  18 static int show_status;
  19
  20 struct stat_rec {
  21   u64 pos;
  22   u64 len;
  23 };
  24
  25 static struct stat_rec *status;
  26
  27 #define ASORT_PREFIX(x) ss_##x
  28 #define ASORT_KEY_TYPE struct stat_rec
  29 #define ASORT_LT(x,y) x.pos < y.pos
  30 #include <ucw/sorter/array-simple.h>
  31
  32 static void stat_read(int fd)
  33 {
  34   u64 len = ucw_seek(fd, 0, SEEK_END);
  35   ucw_seek(fd, 0, SEEK_SET);
  36   if (len % sizeof(struct stat_rec))
  37     die("Malformed status file");
  38   if (len > ~0U)
  39     die("Status file too long");
  40
  41   uns items = len / sizeof(struct stat_rec);
  42   GARY_INIT(status, items);
  43   if (careful_read(fd, status, len) <= 0)
  44     die("Error reading status file");
  45
  46   ss_sort(status, items);
  47
  48 #if 0
  49   for (uint i=0; i<items; i++)
  50     printf("%jd %jd\n", (intmax_t) status[i].pos, (intmax_t) status[i].len);
  51 #endif
  52
  53   for (uint i=1; i<items; i++)
  54     if (status[i].pos < status[i-1].pos + status[i-1].len)
  55       die("Overlapping status file entries: [%jd,%jd] with [%jd,%jd]",
  56         (intmax_t) status[i-1].pos,
  57         (intmax_t)(status[i-1].pos + status[i-1].len),
  58         (intmax_t) status[i].pos,
  59         (intmax_t)(status[i].pos + status[i].len));
  60 }
  61
  62 static u64 stat_have(void)
  63 {
  64   u64 have = 0;
  65   for (uint i=0; i<GARY_SIZE(status); i++)
  66     {
  67       u64 x = CLAMP(status[i].pos, start_pos, end_pos);
  68       u64 y = CLAMP(status[i].pos + status[i].len, start_pos, end_pos);
  69       have += y - x;
  70     }
  71   return have;
  72 }
  73
  74 static void stat_show(void)
  75 {
  76   u64 last = 0;
  77   u64 total_missing = 0;
  78   for (uint i=0; i<GARY_SIZE(status); i++)
  79     {
  80       if (status[i].pos > last)
  81         {
  82           u64 miss = status[i].pos - last;
  83           printf("Missing: %ju +%ju\n", (intmax_t) last, (intmax_t) miss);
  84           total_missing += miss;
  85         }
  86       last = status[i].pos + status[i].len;
  87     }
  88   printf("End at %ju\n", (intmax_t) last);
  89   printf("Total missing: %ju\n", (intmax_t) total_missing);
  90 }
  91
  92 static struct opt_section options = {
  93   OPT_ITEMS {
  94     OPT_HELP("Usage: ddigger [options] block-device output-file status-file"),
  95     OPT_HELP("   or: ddigger [options] --status - - status-file"),
  96     OPT_HELP(""),
  97     OPT_HELP("Options:"),
  98     OPT_HELP_OPTION,
  99     OPT_STRING(OPT_POSITIONAL(1), NULL, dev_name, OPT_REQUIRED, ""),
 100     OPT_STRING(OPT_POSITIONAL(2), NULL, out_name, OPT_REQUIRED, ""),
 101     OPT_STRING(OPT_POSITIONAL(3), NULL, stat_name, OPT_REQUIRED, ""),
 102     OPT_UINT('b', "block-size", block_size, OPT_REQUIRED_VALUE, "<bytes>\tTransfer block size (default=1M)"),
 103     OPT_UINT('e', "error-skip", error_skip, OPT_REQUIRED_VALUE, "<bytes>\tHow far to skip on error (default=block size)"),
 104     OPT_U64(0, "start", start_pos, OPT_REQUIRED_VALUE, "<bytes>\tStart position (default=beginning of device)"),
 105     OPT_U64(0, "end", end_pos, OPT_REQUIRED_VALUE, "<bytes>\tEnd position (default=end of device)"),
 106     OPT_SWITCH(0, "status", show_status, 1, 0, "\tShow contents of status file"),
 107     OPT_END
 108   }
 109 };
 110
 111 int main(int argc UNUSED, char **argv)
 112 {
 113   opt_parse(&options, argv+1);
 114   if (!error_skip)
 115     error_skip = block_size;
 116
 117   int stat_fd = ucw_open(stat_name, (show_status ? O_RDONLY : O_RDWR | O_CREAT), 0666);
 118   if (stat_fd < 0)
 119     die("Cannot open %s: %m", stat_name);
 120   stat_read(stat_fd);
 121
 122   if (show_status)
 123     {
 124       stat_show();
 125       return 0;
 126     }
 127
 128   int dev_fd = ucw_open(dev_name, O_RDONLY | O_DIRECT);
 129   if (dev_fd < 0)
 130     die("Cannot open block device %s: %m", dev_name);
 131
 132   u64 dev_size;
 133   if (ioctl(dev_fd, BLKGETSIZE64, &dev_size) < 0)
 134     die("BLKGETSIZE64: %m");
 135   start_pos = MIN(dev_size, start_pos);
 136   end_pos = MIN(dev_size, end_pos);
 137
 138   int out_fd = ucw_open(out_name, O_RDWR | O_CREAT, 0666);
 139   if (out_fd < 0)
 140     die("Cannot open %s: %m", out_name);
 141
 142   if (ucw_ftruncate(out_fd, dev_size) < 0)
 143     die("Cannot resize %s: %m", out_name);
 144
 145   u64 remains = (end_pos - start_pos) + stat_have();
 146   ucw_seek(stat_fd, 0, SEEK_END);
 147   *GARY_PUSH(status) = (struct stat_rec) { .pos = end_pos, .len = 0 };
 148
 149   byte *buf = big_alloc(MAX(block_size, 4096));
 150   u64 pos = start_pos;
 151   u64 have = 0;
 152   uint stat_i = 0;
 153   uint errors = 0;
 154
 155   while (pos < end_pos)
 156     {
 157       while (status[stat_i].pos + status[stat_i].len < pos)
 158         stat_i++;
 159       if (status[stat_i].pos <= pos)
 160         {
 161           pos = status[stat_i].pos + status[stat_i].len;
 162           stat_i++;
 163           continue;
 164         }
 165
 166       u64 len = MIN(block_size, status[stat_i].pos - pos);
 167       if (pos % block_size)
 168         len = MIN(len, block_size - pos % block_size);
 169
 170 #define GAUGE(curr,max) (curr), (max), 100*(double)(curr)/(double)((max)?:1)
 171       fprintf(stderr, "Pos: %zd / %zd (%.2f%%)  Rel: %zd / %zd (%.2f%%)  Have: %zd / %zd (%.2f%%)  Errors: %u\r",
 172         GAUGE(pos, dev_size),
 173         GAUGE(pos-start_pos, end_pos-start_pos),
 174         GAUGE(have, remains),
 175         errors);
 176 #undef GAUGE
 177
 178       if (ucw_seek(dev_fd, pos, SEEK_SET) < 0)
 179         die("lseek: %m");
 180       if (ucw_seek(out_fd, pos, SEEK_SET) < 0)
 181         die("lseek: %m");
 182
 183       ssize_t done = read(dev_fd, buf, len);
 184       if (done < (ssize_t) len)
 185         {
 186           fprintf(stderr, "\n");
 187           // msg(L_INFO, "Error: %d bytes read at position %jd", (int) MAX(done, 0), (intmax_t) pos);
 188           errors++;
 189           pos += error_skip;
 190         }
 191       else
 192         {
 193           if (careful_write(out_fd, buf, len) < 0)
 194             die("Error writing %s: %m", out_name);
 195
 196           struct stat_rec sr = { .pos = pos, .len = len };
 197           if (careful_write(stat_fd, &sr, sizeof(sr)) < 0)
 198             die("Error writing status file: %m");
 199
 200           pos += len;
 201           have += len;
 202         }
 203     }
 204
 205   return 0;
 206 }