mirror of
https://github.com/upa/mscp.git
synced 2026-02-13 00:24:42 +08:00
copy over multiple connection works, but not checked file integirty
This commit is contained in:
40
src/file.c
40
src/file.c
@@ -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);
|
||||
|
||||
246
src/main.c
246
src/main.c
@@ -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
|
||||
|
||||
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);
|
||||
/* 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;
|
||||
}
|
||||
|
||||
|
||||
ssh_sftp_close(sscp.ctrl);
|
||||
return 0;
|
||||
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;
|
||||
|
||||
/* 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;
|
||||
|
||||
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" \
|
||||
|
||||
Reference in New Issue
Block a user