ann.py: Further improvements.

[utils.git] / pdu.c

#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <unistd.h>
#include <errno.h>
#include <dirent.h>
#include <stdarg.h>
#include <stdint.h>
#include <sys/stat.h>
#include <dpar.h>
#include <stdatomic.h>

struct subdir {
    struct subdir *parent;
    char *path;
    off_t sum;
    int rv, subjobs, summary;
};

static int summary = 0;
static int retval = 0;

static char *sprintf2(char *fmt, ...)
{
    char *ret;
    size_t sz;
    FILE *fp;
    va_list args;
    
    fp = open_memstream(&ret, &sz);
    va_start(args, fmt);
    vfprintf(fp, fmt, args);
    va_end(args);
    fclose(fp);
    return(ret);
}

static void printres(struct subdir *job)
{
    printf("%ji\t%s\n", (intmax_t)(job->sum / 1024), job->path);
}

static void finish(struct subdir *job)
{
    if(atomic_fetch_sub(&job->subjobs, 1) == 1) {
	if(job->summary)
	    printres(job);
	if(job->parent) {
	    atomic_fetch_add(&job->parent->sum, job->sum);
	    finish(job->parent);
	}
	free(job->path);
	free(job);
    }
}

static void dodir(void *jobp)
{
    struct subdir *job = jobp;
    DIR *dp;
    struct dirent *dent, dbuf;
    struct stat sb;
    struct subdir *nj;
    char *ep;
    
    if((dp = opendir(job->path)) == NULL) {
	fprintf(stderr, "pdu: %s: %s\n", job->path, strerror(errno));
	job->rv = 1;
	return;
    }
    job->subjobs = 1;
    while(1) {
	if(readdir_r(dp, &dbuf, &dent)) {
	    fprintf(stderr, "pdu: reading %s: %s\n", job->path, strerror(errno));
	    job->rv = 1;
	    break;
	}
	if(!dent)
	    break;
	if(!strcmp(dent->d_name, ".") || !strcmp(dent->d_name, ".."))
	    continue;
	ep = sprintf2("%s/%s", job->path, dent->d_name);
	if(stat(ep, &sb)) {
	    fprintf(stderr, "pdu: stat %s: %s\n", ep, strerror(errno));
	} else {
	    if(S_ISDIR(sb.st_mode)) {
		*(nj = malloc(sizeof(*nj))) = (struct subdir) {
		    .path = strdup(ep),
		    .summary = !summary,
		    .parent = job,
		};
		atomic_fetch_add(&job->subjobs, 1);
		atomic_fetch_add(&job->sum, (off_t)sb.st_blocks * 512);
		submitiowork(dodir, nj);
	    } else if(S_ISREG(sb.st_mode)) {
		atomic_fetch_add(&job->sum, (off_t)sb.st_blocks * 512);
	    }
	}
	free(ep);
    }
    closedir(dp);
    finish(job);
}

static void dodu(char *path)
{
    struct subdir *job;
    
    *(job = malloc(sizeof(*job))) = (struct subdir) {
	.path = strdup(path),
	.summary = 1,
    };
    submitiowork(dodir, job);
}

static void usage(FILE *out)
{
    fprintf(out, "usage: pdu [-sh] [DIRECTORY...]\n");
}

int main(int argc, char **argv)
{
    int i, c;
    
    while((c = getopt(argc, argv, "hs")) != -1) {
	switch(c) {
	case 's':
	    summary = 1;
	    break;
	case 'h':
	    usage(stdout);
	    exit(0);
	default:
	    usage(stderr);
	    exit(1);
	}
    }
    if(optind < argc) {
	for(i = optind; i < argc; i++)
	    dodu(argv[i]);
    } else {
	dodu(".");
    }
    joinwork();
    return(retval);
}

/*
 * Local Variables:
 * compile-command: "gcc -Wall -g -o pdu pdu.c -ldpar"
 * End:
 */