sys/src/cmd/pump.c

   1 /* pump - copy through circular buffer */
   2 #include <u.h>
   3 #include <libc.h>
   4
   5 uchar*  buf;
   6
   7 Lock    arithlock;      /* protect 64-bit accesses: unlikely to be atomic */
   8 uvlong  nin;
   9 uvlong  nout;
  10
  11 ulong   kilo;
  12 ulong   max;
  13 long    ssize;
  14 vlong   tsize;
  15 int     dsize;
  16 int     done;
  17 int     ibsize;
  18 int     obsize;
  19 int     verb;
  20
  21 void    doinput(int);
  22 void    dooutput(int);
  23
  24 static void
  25 usage(void)
  26 {
  27         fprint(2, "usage: pump [-f ofile] [-k KB-buffer] [-i ireadsize]\n"
  28                 "\t[-o owritesize] [-b iando] [-s start-KB] [-d sleeptime] "
  29                 "[files]\n");
  30         exits("usage");
  31 }
  32
  33 void
  34 main(int argc, char *argv[])
  35 {
  36         int i, f, fo;
  37         char *file;
  38
  39         kilo = 5000;
  40         obsize = ibsize = 8*1024;
  41         dsize = 0;
  42         fo = 1;
  43
  44         ARGBEGIN {
  45         default:
  46                 usage();
  47         case 'b':
  48                 obsize = ibsize = atoi(EARGF(usage()));
  49                 break;
  50         case 'd':
  51                 dsize = atoi(EARGF(usage()));
  52                 break;
  53         case 'f':
  54                 file = EARGF(usage());
  55                 fo = create(file, 1, 0666);
  56                 if(fo < 0)
  57                         sysfatal("can't create %s: %r", file);
  58                 break;
  59         case 'i':
  60                 ibsize = atoi(EARGF(usage()));
  61                 break;
  62         case 'k':
  63                 kilo = atoi(EARGF(usage()));
  64                 break;
  65         case 'o':
  66                 obsize = atoi(EARGF(usage()));
  67                 break;
  68         case 's':
  69                 ssize = atoi(EARGF(usage()));
  70                 if(ssize <= 0)
  71                         ssize = 800;
  72                 ssize <<= 10;
  73                 break;
  74         case 't':
  75                 tsize = atoll(EARGF(usage()));
  76                 tsize *= 10584000;              /* minutes */
  77                 break;
  78         } ARGEND
  79         kilo <<= 10;
  80
  81         buf = malloc(kilo);
  82         if(buf == nil)
  83                 sysfatal("no memory: %r");
  84         nin = 0;
  85         nout = 0;
  86         done = 0;
  87         max = 0;
  88
  89         switch(rfork(RFPROC|RFNOWAIT|RFNAMEG|RFMEM)) {
  90         default:
  91                 dooutput(fo);
  92                 break;
  93         case 0:
  94                 for(i=0; i<argc; i++) {
  95                         f = open(argv[i], OREAD);
  96                         if(f < 0) {
  97                                 fprint(2, "%s: can't open %s: %r\n",
  98                                         argv0, argv[i]);
  99                                 break;
 100                         }
 101                         doinput(f);
 102                         close(f);
 103                 }
 104                 if(argc == 0)
 105                         doinput(0);
 106                 break;
 107         case -1:
 108                 fprint(2, "%s: fork failed: %r\n", argv0);
 109                 break;
 110         }
 111         done = 1;
 112         exits(0);
 113 }
 114
 115 /* call with arithlock held */
 116 static int
 117 sleepunlocked(long ms)
 118 {
 119         int r;
 120
 121         unlock(&arithlock);
 122         r = sleep(ms);
 123         lock(&arithlock);
 124         return r;
 125 }
 126
 127 void
 128 dooutput(int f)
 129 {
 130         long n, l, c;
 131
 132         lock(&arithlock);
 133         for (;;) {
 134                 n = nin - nout;
 135                 if(n == 0) {
 136                         if(done)
 137                                 break;
 138                         sleepunlocked(dsize);
 139                         continue;
 140                 }
 141                 if(verb && n > max) {
 142                         fprint(2, "n = %ld\n", n);
 143                         max = n;
 144                 }
 145                 l = nout % kilo;
 146                 unlock(&arithlock);
 147
 148                 if(kilo-l < n)
 149                         n = kilo-l;
 150                 if(n > obsize)
 151                         n = obsize;
 152                 c = write(f, buf+l, n);
 153
 154                 lock(&arithlock);
 155                 if(c != n) {
 156                         fprint(2, "%s: write error: %r\n", argv0);
 157                         break;
 158                 }
 159                 nout += c;
 160                 if(tsize && nout > tsize) {
 161                         fprint(2, "%s: time limit exceeded\n", argv0);
 162                         break;
 163                 }
 164         }
 165         unlock(&arithlock);
 166 }
 167
 168 void
 169 doinput(int f)
 170 {
 171         long n, l, c, xnin;
 172
 173         lock(&arithlock);
 174         if(ssize > 0) {
 175                 for (xnin = 0; xnin < ssize && !done; xnin += c) {
 176                         n = kilo - (xnin - nout);
 177                         if(n == 0)
 178                                 break;
 179                         unlock(&arithlock);
 180
 181                         l = xnin % kilo;
 182                         if(kilo-l < n)
 183                                 n = kilo-l;
 184                         if(n > ibsize)
 185                                 n = ibsize;
 186                         c = read(f, buf+l, n);
 187
 188                         lock(&arithlock);
 189                         if(c <= 0) {
 190                                 if(c < 0)
 191                                         fprint(2, "%s: read error: %r\n", argv0);
 192                                 break;
 193                         }
 194                 }
 195                 nin = xnin;
 196         }
 197         while(!done) {
 198                 n = kilo - (nin - nout);
 199                 if(n == 0) {
 200                         sleepunlocked(0);
 201                         continue;
 202                 }
 203                 l = nin % kilo;
 204                 unlock(&arithlock);
 205
 206                 if(kilo-l < n)
 207                         n = kilo-l;
 208                 if(n > ibsize)
 209                         n = ibsize;
 210                 c = read(f, buf+l, n);
 211
 212                 lock(&arithlock);
 213                 if(c <= 0) {
 214                         if(c < 0)
 215                                 fprint(2, "%s: read error: %r\n", argv0);
 216                         break;
 217                 }
 218                 nin += c;
 219         }
 220         unlock(&arithlock);
 221 }