Subversion Repositories planix.SVN

/* pump - copy through circular buffer */
#include <u.h>
#include <libc.h>

uchar*  buf;

Lock    arithlock;      /* protect 64-bit accesses: unlikely to be atomic */
uvlong  nin;
uvlong  nout;

ulong   kilo;
ulong   max;
long    ssize;
vlong   tsize;
int     dsize;
int     done;
int     ibsize;
int     obsize;
int     verb;
vlong   off;

void    doinput(int);
void    dooutput(int);

static void
usage(void)
{
        fprint(2, "usage: pump [-b iando] [-d sleeptime] [-f ofile] "
                "[-i ireadsize]\n\t[-k KB-buffer] [-o owritesize] "
                "[-s start-KB] [-S seek-offset]\n\t[-t mins] [files]\n");
        exits("usage");
}

void
main(int argc, char *argv[])
{
        int i, f, fo;
        char *file;

        kilo = 5000;
        obsize = ibsize = 8*1024;
        dsize = 0;
        fo = 1;
        off = 0;

        ARGBEGIN {
        default:
                usage();
        case 'b':
                obsize = ibsize = atoi(EARGF(usage()));
                break;
        case 'd':
                dsize = atoi(EARGF(usage()));
                break;
        case 'f':
                file = EARGF(usage());
                fo = create(file, 1, 0666);
                if(fo < 0)
                        sysfatal("can't create %s: %r", file);
                break;
        case 'i':
                ibsize = atoi(EARGF(usage()));
                break;
        case 'k':
                kilo = atoi(EARGF(usage()));
                break;
        case 'o':
                obsize = atoi(EARGF(usage()));
                break;
        case 's':
                ssize = atoi(EARGF(usage()));
                if(ssize <= 0)
                        ssize = 800;
                ssize <<= 10;
                break;
        case 'S':
                off = atoll(EARGF(usage()));
                if(off < 0)
                        sysfatal("seek offset %lld must be non-negative", off);
                break;
        case 't':
                tsize = atoll(EARGF(usage()));
                tsize *= 10584000;              /* minutes */
                break;
        } ARGEND
        kilo <<= 10;

        buf = malloc(kilo);
        if(buf == nil)
                sysfatal("no memory: %r");
        nin = 0;
        nout = 0;
        done = 0;
        max = 0;

        switch(rfork(RFPROC|RFNOWAIT|RFNAMEG|RFMEM)) {
        default:
                dooutput(fo);
                break;
        case 0:
                for(i=0; i<argc; i++) {
                        f = open(argv[i], OREAD);
                        if(f < 0) {
                                fprint(2, "%s: can't open %s: %r\n",
                                        argv0, argv[i]);
                                break;
                        }
                        doinput(f);
                        close(f);
                }
                if(argc == 0)
                        doinput(0);
                break;
        case -1:
                fprint(2, "%s: fork failed: %r\n", argv0);
                break;
        }
        done = 1;
        exits(0);
}

/* call with arithlock held */
static int
sleepunlocked(long ms)
{
        int r;

        unlock(&arithlock);
        r = sleep(ms);
        lock(&arithlock);
        return r;
}

void
dooutput(int f)
{
        long n, l, c;

        seek(f, off, 0);
        lock(&arithlock);
        for (;;) {
                n = nin - nout;
                if(n == 0) {
                        if(done)
                                break;
                        sleepunlocked(dsize);
                        continue;
                }
                if(verb && n > max) {
                        fprint(2, "n = %ld\n", n);
                        max = n;
                }
                l = nout % kilo;
                unlock(&arithlock);

                if(kilo-l < n)
                        n = kilo-l;
                if(n > obsize)
                        n = obsize;
                c = write(f, buf+l, n);

                lock(&arithlock);
                if(c != n) {
                        fprint(2, "%s: write error at offset %,lld: %r\n",
                                argv0, seek(f, 0, 1));
                        break;
                }
                nout += c;
                if(tsize && nout > tsize) {
                        fprint(2, "%s: time limit exceeded\n", argv0);
                        break;
                }
        }
        unlock(&arithlock);
}

void
doinput(int f)
{
        long n, l, c, xnin;

        seek(f, off, 0);
        lock(&arithlock);
        if(ssize > 0) {
                for (xnin = 0; xnin < ssize && !done; xnin += c) {
                        n = kilo - (xnin - nout);
                        if(n == 0)
                                break;
                        unlock(&arithlock);

                        l = xnin % kilo;
                        if(kilo-l < n)
                                n = kilo-l;
                        if(n > ibsize)
                                n = ibsize;
                        c = read(f, buf+l, n);

                        lock(&arithlock);
                        if(c <= 0) {
                                if(c < 0)
                                        fprint(2, "%s: read error: %r\n", argv0);
                                break;
                        }
                }
                nin = xnin;
        }
        while(!done) {
                n = kilo - (nin - nout);
                if(n == 0) {
                        sleepunlocked(0);
                        continue;
                }
                l = nin % kilo;
                unlock(&arithlock);

                if(kilo-l < n)
                        n = kilo-l;
                if(n > ibsize)
                        n = ibsize;
                c = read(f, buf+l, n);

                lock(&arithlock);
                if(c <= 0) {
                        if(c < 0)
                                fprint(2, "%s: read error: %r\n", argv0);
                        break;
                }
                nin += c;
        }
        unlock(&arithlock);
}