xref: /plan9/sys/src/cmd/pump.c (revision 58da3067adcdccaaa043d0bfde28ba83b7ced07d)
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 at offset %,lld: %r\n",
157 				argv0, seek(f, 0, 1));
158 			break;
159 		}
160 		nout += c;
161 		if(tsize && nout > tsize) {
162 			fprint(2, "%s: time limit exceeded\n", argv0);
163 			break;
164 		}
165 	}
166 	unlock(&arithlock);
167 }
168 
169 void
170 doinput(int f)
171 {
172 	long n, l, c, xnin;
173 
174 	lock(&arithlock);
175 	if(ssize > 0) {
176 		for (xnin = 0; xnin < ssize && !done; xnin += c) {
177 			n = kilo - (xnin - nout);
178 			if(n == 0)
179 				break;
180 			unlock(&arithlock);
181 
182 			l = xnin % kilo;
183 			if(kilo-l < n)
184 				n = kilo-l;
185 			if(n > ibsize)
186 				n = ibsize;
187 			c = read(f, buf+l, n);
188 
189 			lock(&arithlock);
190 			if(c <= 0) {
191 				if(c < 0)
192 					fprint(2, "%s: read error: %r\n", argv0);
193 				break;
194 			}
195 		}
196 		nin = xnin;
197 	}
198 	while(!done) {
199 		n = kilo - (nin - nout);
200 		if(n == 0) {
201 			sleepunlocked(0);
202 			continue;
203 		}
204 		l = nin % kilo;
205 		unlock(&arithlock);
206 
207 		if(kilo-l < n)
208 			n = kilo-l;
209 		if(n > ibsize)
210 			n = ibsize;
211 		c = read(f, buf+l, n);
212 
213 		lock(&arithlock);
214 		if(c <= 0) {
215 			if(c < 0)
216 				fprint(2, "%s: read error: %r\n", argv0);
217 			break;
218 		}
219 		nin += c;
220 	}
221 	unlock(&arithlock);
222 }
223