diff options
| -rw-r--r-- | src/file.c | 40 | ||||
| -rw-r--r-- | src/file.h | 4 | ||||
| -rw-r--r-- | src/main.c | 244 | ||||
| -rw-r--r-- | src/util.h | 4 |
4 files changed, 248 insertions, 44 deletions
@@ -340,13 +340,14 @@ static int file_dst_prepare(struct file *f, sftp_session sftp) if (sftp) { ret = sftp_mkdir(sftp, path, mode); - if (ret < 0) { + if (ret < 0 && + sftp_get_error(sftp) != SSH_FX_FILE_ALREADY_EXISTS) { pr_err("failed to create %s: %s\n", path, ssh_get_error(sftp_ssh(sftp))); return -1; } } else { - if (mkdir(path, mode) == -1) { + if (mkdir(path, mode) == -1 && errno != EEXIST) { pr_err("failed to create %s: %s\n", path, strerrno()); return -1; @@ -487,6 +488,7 @@ int chunk_prepare(struct chunk *c, sftp_session sftp) goto out; } f->state = FILE_STATE_OPENED; + pr("copy start: %s\n", f->path); } out: @@ -577,11 +579,12 @@ static sftp_file chunk_open_remote(const char *path, int flags, size_t off, return sf; } -static int chunk_copy_local_to_remote(struct chunk *c, sftp_session sftp, size_t buf_sz) +static int chunk_copy_local_to_remote(struct chunk *c, sftp_session sftp, size_t buf_sz, + size_t *counter) { struct file *f = c->f; char buf[buf_sz]; - size_t remaind, remaind2; + size_t remaind, remaind2, read_size; sftp_file sf = NULL; mode_t mode; int fd = 0; @@ -598,12 +601,15 @@ static int chunk_copy_local_to_remote(struct chunk *c, sftp_session sftp, size_t } for (remaind = c->len; remaind > 0;) { - ret = read(fd, buf, buf_sz); + read_size = buf_sz < remaind ? buf_sz : remaind; + ret = read(fd, buf, read_size); if (ret < 0) { pr_err("failed to read %s: %s\n", f->path, strerrno()); ret = -1; goto out; } + if (ret == 0) + break; for (remaind2 = ret; remaind2 > 0;) { ret2 = sftp_write(sf, buf + (ret - remaind2), remaind2); @@ -613,8 +619,9 @@ static int chunk_copy_local_to_remote(struct chunk *c, sftp_session sftp, size_t ret = -1; goto out; } - remaind2 -= ret2; c->done += ret2; + *counter += ret2; + remaind2 -= ret2; } remaind -= ret; @@ -636,11 +643,12 @@ out: return ret; } -static int chunk_copy_remote_to_local(struct chunk *c, sftp_session sftp, size_t buf_sz) +static int chunk_copy_remote_to_local(struct chunk *c, sftp_session sftp, size_t buf_sz, + size_t *counter) { struct file *f = c->f; char buf[buf_sz]; - size_t remaind, remaind2; + size_t remaind, remaind2, read_size; sftp_file sf = NULL; mode_t mode; int fd = 0; @@ -657,7 +665,8 @@ static int chunk_copy_remote_to_local(struct chunk *c, sftp_session sftp, size_t } for (remaind = c->len; remaind > 0;) { - ret = sftp_read(sf, buf, buf_sz); + read_size = buf_sz < remaind ? buf_sz : remaind; + ret = sftp_read(sf, buf, read_size); if (ret < 0) { pr_err("failed to read from %s: %s\n", f->dst_path, ssh_get_error(sftp_ssh(sftp))); @@ -673,8 +682,9 @@ static int chunk_copy_remote_to_local(struct chunk *c, sftp_session sftp, size_t ret = -1; goto out; } - remaind2 -= ret2; c->done += ret2; + *counter += ret2; + remaind2 -= ret2; } remaind -= ret; @@ -697,19 +707,19 @@ out: return ret; } -int chunk_copy(struct chunk *c, sftp_session sftp, size_t buf_sz) +int chunk_copy(struct chunk *c, sftp_session sftp, size_t buf_sz, size_t *counter) { struct file *f = c->f; int ret; - pr_debug("copy %s %s -> %s %s\n", + pr_debug("copy %s %s -> %s %s off=0x%010lx\n", f->path, f->remote ? "(remote)" : "(local)", - f->dst_path, f->dst_remote ? "(remote)" : "(local)") + f->dst_path, f->dst_remote ? "(remote)" : "(local)", c->off); if (f->dst_remote) - ret = chunk_copy_local_to_remote(c, sftp, buf_sz); + ret = chunk_copy_local_to_remote(c, sftp, buf_sz, counter); else - ret = chunk_copy_remote_to_local(c, sftp, buf_sz); + ret = chunk_copy_remote_to_local(c, sftp, buf_sz, counter); if (ret < 0) return ret; @@ -22,8 +22,6 @@ struct file { int state; /* destination file state */ lock lock; /* mutex to protect state */ refcnt refcnt; /* chunks referencing this file */ - - size_t done; /* copied bytes. a control thread totaling up done of chunks */ }; #define FILE_STATE_INIT 0 #define FILE_STATE_OPENED 1 @@ -72,7 +70,7 @@ int chunk_fill(struct list_head *file_list, struct list_head *chunk_list, struct chunk *chunk_acquire(struct list_head *chunk_list); int chunk_prepare(struct chunk *c, sftp_session sftp); -int chunk_copy(struct chunk *c, sftp_session sftp, size_t buf_sz); +int chunk_copy(struct chunk *c, sftp_session sftp, size_t buf_sz, size_t *counter); #ifdef DEBUG void file_dump(struct list_head *file_list); @@ -2,6 +2,10 @@ #include <stdlib.h> #include <stdbool.h> #include <unistd.h> +#include <signal.h> +#include <sys/time.h> +#include <math.h> +#include <pthread.h> #include <list.h> #include <util.h> @@ -12,16 +16,19 @@ int verbose = 0; /* util.h */ + #define DEFAULT_MIN_CHUNK_SZ (64 << 20) /* 64MB */ #define DEFAULT_BUF_SZ 32768 /* CHANNEL_MAX_PACKET in libssh */ /* XXX: passing over CHANNEL_MAX_PACKET bytes to sftp_write stalls */ struct sscp { char *host; /* remote host (and username) */ + struct ssh_opts *opts; /* ssh parameters */ sftp_session ctrl; /* control sftp session */ + struct list_head file_list; - struct list_head chunk_list; + struct list_head chunk_list; /* stack of chunks */ lock chunk_lock; /* lock for chunk list */ char *target; @@ -29,6 +36,34 @@ struct sscp { int buf_sz; }; +struct sscp_thread { + struct sscp *sscp; + sftp_session sftp; + + pthread_t tid; + size_t done; /* copied bytes */ + bool finished; +}; + +void *sscp_copy_thread(void *arg); +void *sscp_monitor_thread(void *arg); + +static pthread_t mtid; +struct sscp_thread *threads; +int nr_threads; + +void stop_all(int sig) +{ + int n; + + pr("stopping...\n"); + for (n = 0; n < nr_threads; n++) { + pthread_cancel(threads[n].tid); + } + pthread_cancel(mtid); +} + + void usage(bool print_help) { printf("sscp: super scp, copy files over multiple ssh connections\n" "\n" @@ -111,10 +146,9 @@ int main(int argc, char **argv) { struct sscp sscp; struct ssh_opts opts; - int nr_conn = nr_cpus(); int min_chunk_sz = DEFAULT_MIN_CHUNK_SZ; int max_chunk_sz = 0; - int ret = 0; + int ret = 0, n; char ch; memset(&opts, 0, sizeof(opts)); @@ -124,11 +158,13 @@ int main(int argc, char **argv) lock_init(&sscp.chunk_lock); sscp.buf_sz = DEFAULT_BUF_SZ; + nr_threads = nr_cpus(); + while ((ch = getopt(argc, argv, "n:s:S:b:l:p:i:c:Cvh")) != -1) { switch (ch) { case 'n': - nr_conn = atoi(optarg); - if (nr_conn < 1) { + nr_threads = atoi(optarg); + if (nr_threads < 1) { pr_err("invalid number of connections: %s\n", optarg); return 1; } @@ -221,50 +257,208 @@ int main(int argc, char **argv) sscp.ctrl = ssh_make_sftp_session(sscp.host, &opts); if (!sscp.ctrl) return 1; + sscp.opts = &opts; /* save ssh-able ssh_opts */ /* check target is directory */ ret = file_is_directory(sscp.target, file_find_hostname(sscp.target) ? sscp.ctrl : NULL); if (ret < 0) - return 1; + goto out; if (ret == 0) { pr_err("target must be directory\n"); - return 1; + goto out; } /* fill file list */ ret = file_fill(sscp.ctrl, &sscp.file_list, &argv[optind], argc - optind - 1); - if (ret < 0) { - ssh_sftp_close(sscp.ctrl); - return 1; - } + if (ret < 0) + goto out; + ret = file_fill_dst(sscp.target, &sscp.file_list); - if (ret < 0){ - ssh_sftp_close(sscp.ctrl); - return -1; - } + if (ret < 0) + goto out; + #ifdef DEBUG file_dump(&sscp.file_list); #endif /* fill chunk list */ ret = chunk_fill(&sscp.file_list, &sscp.chunk_list, - nr_conn, min_chunk_sz, max_chunk_sz); - if (ret < 0) { - ssh_sftp_close(sscp.ctrl); - return 1; - } + nr_threads, min_chunk_sz, max_chunk_sz); + if (ret < 0) + goto out; + #ifdef DEBUG chunk_dump(&sscp.chunk_list); #endif + /* register SIGINT to stop thrads */ + if (signal(SIGINT, stop_all) == SIG_ERR) { + pr_err("cannot set signal: %s\n", strerrno()); + ret = 1; + goto out; + } + + /* spawn threads */ + threads = calloc(nr_threads, sizeof(struct sscp_thread)); + memset(threads, 0, nr_threads * sizeof(struct sscp_thread)); + for (n = 0; n < nr_threads; n++) { + struct sscp_thread *t = &threads[n]; + t->sscp = &sscp; + t->finished = false; + ret = pthread_create(&t->tid, NULL, sscp_copy_thread, t); + if (ret < 0) { + pr_err("pthread_create error: %d\n", ret); + stop_all(0); + goto join_out; + } + } + + /* spawn count thread */ + ret = pthread_create(&mtid, NULL, sscp_monitor_thread, &sscp); + if (ret < 0) { + pr_err("pthread_create error: %d\n", ret); + stop_all(0); + goto join_out; + } + + +join_out: + /* waiting for threads join... */ + for (n = 0; n < nr_threads; n++) + if (threads[n].tid) + pthread_join(threads[n].tid, NULL); + + if (mtid != 0) + pthread_join(mtid, NULL); + +out: + if (sscp.ctrl) + ssh_sftp_close(sscp.ctrl); + + return ret; +} + +void sscp_copy_thread_cleanup(void *arg) +{ + struct sscp_thread *t = arg; + if (t->sftp) + ssh_sftp_close(t->sftp); + t->finished = true; +} + +void *sscp_copy_thread(void *arg) +{ + struct sscp_thread *t = arg; + struct sscp *sscp = t->sscp; + sftp_session sftp; struct chunk *c; - list_for_each_entry(c, &sscp.chunk_list, list) { - chunk_prepare(c, sscp.ctrl); - chunk_copy(c, sscp.ctrl, sscp.buf_sz); + + /* create sftp session */ + sftp = ssh_make_sftp_session(sscp->host, sscp->opts); + if (!sftp) + return NULL; + + pthread_setcancelstate(PTHREAD_CANCEL_ENABLE, NULL); + pthread_setcanceltype(PTHREAD_CANCEL_DEFERRED, NULL); + pthread_cleanup_push(sscp_copy_thread_cleanup, t); + + while (1) { + lock_acquire(&sscp->chunk_lock); + c = chunk_acquire(&sscp->chunk_list); + lock_release(&sscp->chunk_lock); + + if (!c) + break; /* no more chunks */ + + if (chunk_prepare(c, sftp) < 0) + break; + + if (chunk_copy(c, sftp, sscp->buf_sz, &t->done) < 0) + break; + } + + pthread_cleanup_pop(1); + + return NULL; +} + +static double calculate_bps(size_t diff, struct timeval *b, struct timeval *a) +{ + double sec, usec; + + if (a->tv_usec < b->tv_usec) { + a->tv_usec += 1000000; + a->tv_sec--; } + sec = a->tv_sec - b->tv_sec; + usec = a->tv_usec - b->tv_usec; + sec += usec / 1000000; + + return (double)diff / sec * 8; +} + +void *sscp_monitor_thread(void *arg) +{ + struct sscp *sscp = arg; + struct sscp_thread *t; + struct timeval a, b; + struct file *f; + bool all_done; + size_t total, total_round, done, last; + int percent; + double bps; + char *bps_units[] = { "bps", "Kbps", "Mbps", "Gbps" }; + char *byte_units[] = { "B", "KB", "MB", "GB", "TB" }; + int n, bps_u, byte_tu, byte_du; + + total = 0; + done = 0; + last = 0; + + /* get total byte to be transferred */ + list_for_each_entry(f, &sscp->file_list, list) { + total += f->size; + } + total_round = total; + for (byte_tu = 0; total_round > 1000 && byte_tu < 5; byte_tu++) + total_round /= 1024; + + while (1) { + + gettimeofday(&b, NULL); + sleep(1); + + all_done = true; + done = 0; - ssh_sftp_close(sscp.ctrl); - return 0; + for (n = 0; n < nr_threads; n++) { + t = &threads[n]; + done += t->done; + if (!t->finished) + all_done = false; + } + + gettimeofday(&a, NULL); + + percent = floor(((double)(done) / (double)total) * 100); + for (byte_du = 0; done > 1000 && byte_du < 5; byte_du++) done /= 1024; + + bps = calculate_bps(done - last, &b, &a); + for (bps_u = 0; bps > 1000 && bps_u < 4; bps_u++) bps /= 1000; + + printf("%d%% (%lu%s/%lu%s) %.2f %s\n", + percent, + done, byte_units[byte_du], total_round, byte_units[byte_tu], + bps, bps_units[bps_u]); + + if (all_done || total == done) + break; + + last = done; + b = a; + } + + return NULL; } @@ -23,7 +23,9 @@ extern int verbose; #define pr_v3(fmt, ...) pr_v(3, fmt, ##__VA_ARGS__) -#define pr_info(fmt, ...) fprintf(stdout, "%s(): " fmt, \ +#define pr(fmt, ...) fprintf(stderr, fmt, ##__VA_ARGS__) + +#define pr_info(fmt, ...) fprintf(stderr, "INFO:%s(): " fmt, \ __func__, ##__VA_ARGS__) #define pr_warn(fmt, ...) fprintf(stderr, "\x1b[1m\x1b[33m" \ |