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