vidir: Improve zero padding
[moreutils.git] / sponge.c
blobf852ad5a81b45fef88e53c4a91855d816e348897
1 /*
2 * sponge.c - read in all available info from stdin, then output it to
3 * file named on the command line
5 * Copyright © 2006 Tollef Fog Heen
7 * This program is free software; you can redistribute it and/or
8 * modify it under the terms of the GNU General Public License
9 * version 2 as published by the Free Software Foundation.
11 * This program is distributed in the hope that it will be useful,
12 * but WITHOUT ANY WARRANTY; without even the implied warranty of
13 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
14 * General Public License for more details.
16 * You should have received a copy of the GNU General Public License
17 * along with this program; if not, write to the Free Software
18 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307
19 * USA
23 #include <stdio.h>
24 #include <stdlib.h>
25 #include <unistd.h>
26 #include <sys/types.h>
27 #include <sys/stat.h>
28 /* MAX() */
29 #include <sys/param.h>
30 #include <fcntl.h>
31 #include <errno.h>
32 #include <string.h>
33 #include <sys/resource.h>
34 /* SIZE_MAX */
35 #include <stdint.h>
36 #include <signal.h>
37 #include <getopt.h>
39 #include "physmem.c"
41 #define BUFF_SIZE 8192
42 #define MIN_SPONGE_SIZE BUFF_SIZE
43 char *tmpname = NULL;
45 void usage() {
46 printf("sponge [-a] <file>: soak up all input from stdin and write it "
47 "to <file>\n");
48 exit(0);
51 /* all the signal stuff copied from gnu sort */
53 /* The set of signals that are caught. */
54 static sigset_t caught_signals;
56 /* Critical section status. */
57 struct cs_status {
58 int valid; // was bool
59 sigset_t sigs;
62 /* Enter a critical section. */
63 static struct cs_status cs_enter (void) {
64 struct cs_status status;
65 status.valid = (sigprocmask(SIG_BLOCK, &caught_signals, &status.sigs) == 0);
66 return status;
69 /* Leave a critical section. */
70 static void cs_leave (struct cs_status status) {
71 if (status.valid) {
72 /* Ignore failure when restoring the signal mask. */
73 sigprocmask(SIG_SETMASK, &status.sigs, NULL);
77 static void cleanup () {
78 if (tmpname) {
79 unlink(tmpname);
83 static void onexit_cleanup (void) {
84 struct cs_status cs = cs_enter();
85 cleanup();
86 cs_leave(cs);
89 static void sighandler (int sig) {
90 if (! SA_NOCLDSTOP)
91 signal(sig, SIG_IGN);
93 cleanup();
95 signal(sig, SIG_DFL);
96 raise(sig);
99 /* taken from coreutils sort */
100 static size_t default_sponge_size (void) {
101 /* Let MEM be available memory or 1/8 of total memory, whichever
102 is greater. */
103 double avail = physmem_available();
104 double total = physmem_total();
105 double mem = MAX(avail, total / 8);
106 struct rlimit rlimit;
108 /* Let SIZE be MEM, but no more than the maximum object size or
109 system resource limits. Avoid the MIN macro here, as it is not
110 quite right when only one argument is floating point. Don't
111 bother to check for values like RLIM_INFINITY since in practice
112 they are not much less than SIZE_MAX. */
113 size_t size = SIZE_MAX;
114 if (mem < size)
115 size = mem;
116 if (getrlimit(RLIMIT_DATA, &rlimit) == 0 && rlimit.rlim_cur < size)
117 size = rlimit.rlim_cur;
118 #ifdef RLIMIT_AS
119 if (getrlimit(RLIMIT_AS, &rlimit) == 0 && rlimit.rlim_cur < size)
120 size = rlimit.rlim_cur;
121 #endif
123 /* Leave a large safety margin for the above limits, as failure can
124 occur when they are exceeded. */
125 size /= 2;
127 #ifdef RLIMIT_RSS
128 /* Leave a 1/16 margin for RSS to leave room for code, stack, etc.
129 Exceeding RSS is not fatal, but can be quite slow. */
130 if (getrlimit(RLIMIT_RSS, &rlimit) == 0 && rlimit.rlim_cur / 16 * 15 < size)
131 size = rlimit.rlim_cur / 16 * 15;
132 #endif
134 /* Use no less than the minimum. */
135 return MAX (size, MIN_SPONGE_SIZE);
138 void trapsignals (void) {
139 ssize_t i = 0;
140 static int const sig[] = {
141 /* The usual suspects. */
142 SIGALRM, SIGHUP, SIGINT, SIGPIPE, SIGQUIT, SIGTERM,
143 #ifdef SIGPOLL
144 SIGPOLL,
145 #endif
146 #ifdef SIGPROF
147 SIGPROF,
148 #endif
149 #ifdef SIGVTALRM
150 SIGVTALRM,
151 #endif
152 #ifdef SIGXCPU
153 SIGXCPU,
154 #endif
155 #ifdef SIGXFSZ
156 SIGXFSZ,
157 #endif
159 int nsigs = sizeof(sig) / sizeof(sig[0]);
161 #if SA_NOCLDSTOP
162 struct sigaction act;
164 sigemptyset(&caught_signals);
165 for (i = 0; i < nsigs; i++) {
166 sigaction(sig[i], NULL, &act);
167 if (act.sa_handler != SIG_IGN)
168 sigaddset(&caught_signals, sig[i]);
171 act.sa_handler = sighandler;
172 act.sa_mask = caught_signals;
173 act.sa_flags = 0;
175 for (i = 0; i < nsigs; i++)
176 if (sigismember(&caught_signals, sig[i]))
177 sigaction(sig[i], &act, NULL);
178 #else
179 for (i = 0; i < nsigs; i++)
180 if (signal(sig[i], SIG_IGN) != SIG_IGN) {
181 signal(sig[i], sighandler);
182 siginterrupt (sig[i], 1);
184 #endif
187 static void write_buff_tmp(char* buff, size_t length, FILE *fd) {
188 if (fwrite(buff, length, 1, fd) < 1) {
189 perror("error writing buffer to temporary file");
190 fclose(fd);
191 exit(1);
195 static void write_buff_tmp_finish (char* buff, size_t length, FILE *fd) {
196 if (length)
197 write_buff_tmp(buff, length, fd);
198 if (fflush(fd) != 0) {
199 perror("fflush");
200 exit(1);
204 static void write_buff_out (char* buff, size_t length, FILE *fd) {
205 if (fwrite(buff, length, 1, fd) < 1) {
206 perror("error writing buffer to output file");
207 fclose(fd);
208 exit(1);
212 static void copy_file (FILE *infile, FILE *outfile, char *buf, size_t size) {
213 ssize_t i;
214 while ((i = read(fileno(infile), buf, size)) > 0) {
215 write_buff_out(buf, i, outfile);
217 if (i == -1) {
218 perror("read file");
219 fclose(infile);
220 exit(1);
224 static void copy_tmpfile (FILE *tmpfile, FILE *outfile, char *buf, size_t size) {
225 if (lseek(fileno(tmpfile), 0, SEEK_SET)) {
226 perror("could to seek to start of file");
227 fclose(tmpfile);
228 exit(1);
230 copy_file(tmpfile, outfile, buf, size);
231 if (fclose(tmpfile) != 0) {
232 perror("read temporary file");
233 exit(1);
235 if (fclose(outfile) != 0) {
236 perror("error writing buffer to output file");
237 exit(1);
241 FILE *open_tmpfile (void) {
242 struct cs_status cs;
243 int tmpfd;
244 FILE *tmpfile;
245 mode_t mask;
246 char *tmpdir;
247 char const * const template="%s/sponge.XXXXXX";
249 trapsignals();
250 cs = cs_enter();
251 tmpdir = getenv("TMPDIR");
252 if (tmpdir == NULL)
253 tmpdir = "/tmp";
254 /* Subtract 2 for `%s' and add 1 for the trailing NULL. */
255 tmpname=malloc(strlen(tmpdir) + strlen(template) - 2 + 1);
256 if (! tmpname) {
257 perror("failed to allocate memory");
258 exit(1);
260 sprintf(tmpname, template, tmpdir);
261 mask=umask(077);
262 tmpfd = mkstemp(tmpname);
263 umask(mask);
264 atexit(onexit_cleanup); // solaris on_exit(onexit_cleanup, 0);
265 cs_leave(cs);
267 if (tmpfd < 0) {
268 perror("mkstemp failed");
269 exit(1);
271 tmpfile = fdopen(tmpfd, "w+");
272 if (! tmpfile) {
273 perror("fdopen");
274 exit(1);
276 return tmpfile;
279 int main (int argc, char **argv) {
280 char *buf, *bufstart, *outname = NULL;
281 size_t bufsize = BUFF_SIZE;
282 size_t bufused = 0;
283 FILE *outfile, *tmpfile = 0;
284 ssize_t i = 0;
285 size_t mem_available = default_sponge_size();
286 int tmpfile_used=0;
287 int append=0;
288 int opt;
290 while ((opt = getopt(argc, argv, "ha")) != -1) {
291 switch (opt) {
292 case 'h':
293 usage();
294 break;
295 case 'a':
296 append=1;
299 if (optind < argc)
300 outname = argv[optind];
302 tmpfile = open_tmpfile();
303 bufstart = buf = malloc(bufsize);
304 if (!buf) {
305 perror("failed to allocate memory");
306 exit(1);
308 while ((i = read(0, buf, bufsize - bufused)) > 0) {
309 bufused = bufused+i;
310 if (bufused == bufsize) {
311 if ((bufsize*2) >= mem_available) {
312 write_buff_tmp(bufstart, bufused, tmpfile);
313 bufused = 0;
314 tmpfile_used = 1;
316 else {
317 bufsize *= 2;
318 bufstart = realloc(bufstart, bufsize);
319 if (!bufstart) {
320 perror("failed to realloc memory");
321 exit(1);
325 buf = bufstart + bufused;
327 if (i < 0) {
328 perror("failed to read from stdin");
329 exit(1);
332 if (outname) {
333 mode_t mode;
334 struct stat statbuf;
335 int exists = (lstat(outname, &statbuf) == 0);
336 int regfile = exists && S_ISREG(statbuf.st_mode) && ! S_ISLNK(statbuf.st_mode);
338 if (append && regfile) {
339 char *tmpbuf = malloc(bufsize);
340 if (!tmpbuf) {
341 perror("failed to allocate memory");
342 exit(1);
344 outfile = fopen(outname, "r");
345 copy_file(outfile, tmpfile, tmpbuf, bufsize);
346 fclose(outfile);
349 write_buff_tmp_finish(bufstart, bufused, tmpfile);
351 /* Set temp file mode to match either
352 * the old file mode, or the default file
353 * mode for a newly created file. */
354 if (exists) {
355 mode = statbuf.st_mode;
357 else {
358 mode_t mask = umask(0);
359 umask(mask);
360 mode = 0666 & ~mask;
362 if (chmod(tmpname, mode) != 0) {
363 perror("chmod");
364 exit(1);
367 /* If it's a regular file, or does not yet exist,
368 * attempt a fast rename of the temp file. */
369 if ((regfile || ! exists) &&
370 rename(tmpname, outname) == 0) {
371 tmpname=NULL; /* don't try to cleanup tmpname */
373 else {
374 /* Fall back to slow copy. */
375 outfile = fopen(outname, "w");
376 if (!outfile) {
377 perror("error opening output file");
378 exit(1);
380 copy_tmpfile(tmpfile, outfile, bufstart, bufsize);
383 else {
384 if (tmpfile_used) {
385 write_buff_tmp_finish(bufstart, bufused, tmpfile);
386 copy_tmpfile(tmpfile, stdout, bufstart, bufsize);
388 else if (bufused) {
389 /* buffer direct to stdout, no tmpfile */
390 write_buff_out(bufstart, bufused, stdout);
394 return 0;