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 } 222