Skip to content

Commit 4caaaf2

Browse files
upaupa
committed
add -L limit bitrate option (#14)
1 parent 9b8ba69 commit 4caaaf2

File tree

11 files changed

+217
-25
lines changed

11 files changed

+217
-25
lines changed

‎CMakeLists.txt

+2-2
Original file line numberDiff line numberDiff line change
@@ -109,7 +109,7 @@ list(APPEND MSCP_BUILD_INCLUDE_DIRS ${CMAKE_CURRENT_BINARY_DIR}/include)
109109
# libmscp.a
110110
set(LIBMSCP_SRC
111111
src/mscp.c src/ssh.c src/fileops.c src/path.c src/checkpoint.c
112-
src/platform.c src/print.c src/pool.c src/strerrno.c
112+
src/bwlimit.c src/platform.c src/print.c src/pool.c src/strerrno.c
113113
${OPENBSD_COMPAT_SRC})
114114
add_library(mscp-static STATIC ${LIBMSCP_SRC})
115115
target_include_directories(mscp-static
@@ -203,7 +203,7 @@ foreach(x RANGE ${DIST_LISTLEN})
203203
COMMENT "Test mscp in ${DOCKER_IMAGE} container"
204204
WORKING_DIRECTORY ${CMAKE_BINARY_DIR}
205205
COMMAND
206-
${CE} run --init --rm --sysctl net.ipv6.conf.all.disable_ipv6=0
206+
${CE} run --init --rm -it --sysctl net.ipv6.conf.all.disable_ipv6=0
207207
${DOCKER_IMAGE} /mscp/scripts/test-in-container.sh)
208208

209209
list(APPEND DOCKER_BUILDS docker-build-${DOCKER_INDEX})

‎doc/mscp.1.in

+8
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,9 @@ mscp \- copy files over multiple SSH connections
3838
.BI \-b \ BUF_SIZE\c
3939
]
4040
[\c
41+
.BI \-L \ LIMIT_BITRATE\c
42+
]
43+
[\c
4144
.BI \-l \ LOGIN_NAME\c
4245
]
4346
[\c
@@ -198,6 +201,11 @@ Specifies the buffer size for I/O and transfer over SFTP. The default
198201
value is 16384. Note that the SSH specification restricts buffer size
199202
delivered over SSH. Changing this value is not recommended at present.
200203

204+
.TP
205+
.B \-L \fILIMIT_BITRATE\fR
206+
Limits the bitrate, specified with k (K), m (M), and g (G), e.g., 100m
207+
indicates 100 Mbps.
208+
201209
.TP
202210
.B \-4
203211
Uses IPv4 addresses only.

‎doc/mscp.rst

+13-9
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
MSCP
33
====
44

5-
:Date: v0.1.4-28-g0d248c5
5+
:Date: v0.1.5-4-g9b8ba69
66

77
NAME
88
====
@@ -12,14 +12,14 @@ mscp - copy files over multiple SSH connections
1212
SYNOPSIS
1313
========
1414

15-
**mscp** [**-46vqDpHdNh**] [ **-n**\ *NR_CONNECTIONS* ] [
16-
**-m**\ *COREMASK* ] [ **-u**\ *MAX_STARTUPS* ] [ **-I**\ *INTERVAL* ] [
17-
**-W**\ *CHECKPOINT* ] [ **-R**\ *CHECKPOINT* ] [
18-
**-s**\ *MIN_CHUNK_SIZE* ] [ **-S**\ *MAX_CHUNK_SIZE* ] [
19-
**-a**\ *NR_AHEAD* ] [ **-b**\ *BUF_SIZE* ] [ **-l**\ *LOGIN_NAME* ] [
20-
**-P**\ *PORT* ] [ **-F**\ *CONFIG* ] [ **-i**\ *IDENTITY* ] [
21-
**-c**\ *CIPHER* ] [ **-M**\ *HMAC* ] [ **-C**\ *COMPRESS* ] [
22-
**-g**\ *CONGESTION* ] *source ... target*
15+
**mscp** [**-46vqDpHdNh**] [ **-n** *NR_CONNECTIONS* ] [ **-m**
16+
*COREMASK* ] [ **-u** *MAX_STARTUPS* ] [ **-I** *INTERVAL* ] [ **-W**
17+
*CHECKPOINT* ] [ **-R** *CHECKPOINT* ] [ **-s** *MIN_CHUNK_SIZE* ] [
18+
**-S** *MAX_CHUNK_SIZE* ] [ **-a** *NR_AHEAD* ] [ **-b** *BUF_SIZE* ] [
19+
**-L** *LIMIT_BITRATE* ] [ **-l** *LOGIN_NAME* ] [ **-P** *PORT* ] [
20+
**-F** *CONFIG* ] [ **-i** *IDENTITY* ] [ **-c** *CIPHER* ] [ **-M**
21+
*HMAC* ] [ **-C** *COMPRESS* ] [ **-g** *CONGESTION* ] *source ...
22+
target*
2323

2424
DESCRIPTION
2525
===========
@@ -111,6 +111,10 @@ OPTIONS
111111
delivered over SSH. Changing this value is not recommended at
112112
present.
113113

114+
**-L LIMIT_BITRATE**
115+
Limits the bitrate, specified with k (K), m (M), and g (G), e.g.,
116+
100m indicates 100 Mbps.
117+
114118
**-4**
115119
Uses IPv4 addresses only.
116120

‎include/mscp.h

+2-1
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,8 @@ struct mscp_opts {
4242
size_t min_chunk_sz; /** minimum chunk size (default 64MB) */
4343
size_t max_chunk_sz; /** maximum chunk size (default file size/nr_threads) */
4444
size_t buf_sz; /** buffer size, default 16k. */
45-
char *coremask; /** hex to specifiy usable cpu cores */
45+
size_t bitrate; /** bits-per-seconds to limit bandwidth */
46+
char *coremask; /** hex to specifiy usable cpu cores */
4647
int max_startups; /** sshd MaxStartups concurrent connections */
4748
int interval; /** interval between SSH connection attempts */
4849
bool preserve_ts; /** preserve file timestamps */

‎src/bwlimit.c

+94
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,94 @@
1+
/* SPDX-License-Identifier: GPL-3.0-only */
2+
#include <errno.h>
3+
4+
#include <bwlimit.h>
5+
#include <platform.h>
6+
7+
#define timespeczerorize(ts) \
8+
do { \
9+
ts.tv_sec = 0; \
10+
ts.tv_nsec = 0; \
11+
} while (0)
12+
13+
int bwlimit_init(struct bwlimit *bw, uint64_t bps, uint64_t win)
14+
{
15+
if (!(bw->sem = sem_create(1)))
16+
return -1;
17+
18+
bw->bps = bps;
19+
bw->win = win; /* msec window */
20+
bw->amt = (double)bps / 8 / 1000 * win; /* bytes in a window (msec) */
21+
bw->credit = bw->amt;
22+
timespeczerorize(bw->wstart);
23+
timespeczerorize(bw->wend);
24+
25+
return 0;
26+
}
27+
28+
#define timespecisset(ts) ((ts).tv_sec || (ts).tv_nsec)
29+
30+
#define timespecmsadd(a, msec, r) \
31+
do { \
32+
(r).tv_sec = (a).tv_sec; \
33+
(r).tv_nsec = (a).tv_nsec + (msec * 1000000); \
34+
if ((r).tv_nsec > 1000000000) { \
35+
(r).tv_sec += (r.tv_nsec) / 1000000000L; \
36+
(r).tv_nsec = (r.tv_nsec) % 1000000000L; \
37+
} \
38+
} while (0)
39+
40+
#define timespecsub(a, b, r) \
41+
do { \
42+
(r).tv_sec = (a).tv_sec - (b).tv_sec; \
43+
(r).tv_nsec = (a).tv_nsec - (b).tv_nsec; \
44+
if ((r).tv_nsec < 0) { \
45+
(r).tv_sec -= 1; \
46+
(r).tv_nsec += 1000000000; \
47+
} \
48+
} while (0)
49+
50+
#define timespeccmp(a, b, expr) \
51+
((a.tv_sec * 1000000000 + a.tv_nsec) expr(b.tv_sec * 1000000000 + b.tv_nsec))
52+
53+
#include <stdio.h>
54+
55+
int bwlimit_wait(struct bwlimit *bw, size_t nr_bytes)
56+
{
57+
struct timespec now, end, rq, rm;
58+
59+
if (bw->bps == 0)
60+
return 0; /* no bandwidth limit */
61+
62+
if (sem_wait(bw->sem) < 0)
63+
return -1;
64+
65+
clock_gettime(CLOCK_MONOTONIC, &now);
66+
67+
if (!timespecisset(bw->wstart)) {
68+
bw->wstart = now;
69+
timespecmsadd(bw->wstart, bw->win, bw->wend);
70+
}
71+
72+
bw->credit -= nr_bytes;
73+
74+
if (bw->credit < 0) {
75+
/* no more credit on this window. sleep until the end
76+
* of this windown and additional time for the
77+
* remaining bytes. */
78+
uint64_t addition = (double)(bw->credit * -1) / (bw->bps / 8);
79+
timespecmsadd(bw->wend, addition * 1000, end);
80+
if (timespeccmp(end, now, >)) {
81+
timespecsub(end, now, rq);
82+
while (nanosleep(&rq, &rm) == -1) {
83+
if (errno != EINTR)
84+
break;
85+
rq = rm;
86+
}
87+
}
88+
bw->credit = bw->amt;
89+
timespeczerorize(bw->wstart);
90+
}
91+
92+
sem_post(bw->sem);
93+
return 0;
94+
}

‎src/bwlimit.h

+27
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
/* SPDX-License-Identifier: GPL-3.0-only */
2+
#ifndef _BWLIMIT_H_
3+
#define _BWLIMIT_H_
4+
5+
#include <stdbool.h>
6+
#include <stdint.h>
7+
#include <time.h>
8+
#include <semaphore.h>
9+
10+
struct bwlimit {
11+
sem_t *sem; /* semaphore */
12+
uint64_t bps; /* limit bit-rate (bps) */
13+
uint64_t win; /* window size (msec) */
14+
size_t amt; /* amount of bytes can be sent in a window */
15+
16+
ssize_t credit; /* remaining bytes can be sent in a window */
17+
struct timespec wstart, wend; /* window start time and end time */
18+
};
19+
20+
int bwlimit_init(struct bwlimit *bw, uint64_t bps, uint64_t win);
21+
/* if bps is 0, it means that bwlimit is not active. If so,
22+
* bwlimit_wait() returns immediately. */
23+
24+
int bwlimit_wait(struct bwlimit *bw, size_t nr_bytes);
25+
26+
27+
#endif /* _BWLIMIT_H_ */

‎src/main.c

+25-2
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,8 @@ void usage(bool print_help)
2626
"\n"
2727
"Usage: mscp [-46vqDpHdNh] [-n nr_conns] [-m coremask]\n"
2828
" [-u max_startups] [-I interval] [-W checkpoint] [-R checkpoint]\n"
29-
" [-s min_chunk_sz] [-S max_chunk_sz] [-a nr_ahead] [-b buf_sz]\n"
29+
" [-s min_chunk_sz] [-S max_chunk_sz] [-a nr_ahead]\n"
30+
" [-b buf_sz] [-L limit_bitrate]\n"
3031
" [-l login_name] [-P port] [-F ssh_config] [-i identity_file]\n"
3132
" [-c cipher_spec] [-M hmac_spec] [-C compress] [-g congestion]\n"
3233
" source ... target\n"
@@ -48,6 +49,7 @@ void usage(bool print_help)
4849
" -S MAX_CHUNK_SIZE max chunk size (default: filesize/nr_conn)\n"
4950
" -a NR_AHEAD number of inflight SFTP commands (default: 32)\n"
5051
" -b BUF_SZ buffer size for i/o and transfer\n"
52+
" -L LIMIT_BITRATE Limit the bitrate, n[KMG] (default: 0)\n"
5153
"\n"
5254
" -4 use IPv4\n"
5355
" -6 use IPv6\n"
@@ -266,12 +268,14 @@ int main(int argc, char **argv)
266268
int direction = 0;
267269
char *remote = NULL, *checkpoint_save = NULL, *checkpoint_load = NULL;
268270
bool dryrun = false, resume = false;
271+
char *u;
272+
size_t mag = 0;
269273

270274
memset(&s, 0, sizeof(s));
271275
memset(&o, 0, sizeof(o));
272276
o.severity = MSCP_SEVERITY_WARN;
273277

274-
#define mscpopts "n:m:u:I:W:R:s:S:a:b:46vqDrl:P:i:F:c:M:C:g:pHdNh"
278+
#define mscpopts "n:m:u:I:W:R:s:S:a:b:L:46vqDrl:P:i:F:c:M:C:g:pHdNh"
275279
while ((ch = getopt(argc, argv, mscpopts)) != -1) {
276280
switch (ch) {
277281
case 'n':
@@ -309,6 +313,25 @@ int main(int argc, char **argv)
309313
case 'b':
310314
o.buf_sz = atoi(optarg);
311315
break;
316+
case 'L':
317+
u = optarg + (strlen(optarg) - 1);
318+
if (*u == 'k' || *u == 'K') {
319+
mag = 1000;
320+
*u = '\0';
321+
} else if (*u == 'm' || *u == 'M') {
322+
mag = 1000000;
323+
*u = '\0';
324+
} else if (*u == 'g' || *u == 'G') {
325+
mag = 1000000000;
326+
*u = '\0';
327+
}
328+
o.bitrate = atol(optarg);
329+
if (o.bitrate == 0) {
330+
pr_err("invalid bitrate: %s", optarg);
331+
return 1;
332+
}
333+
o.bitrate *= mag;
334+
break;
312335
case '4':
313336
s.ai_family = AF_INET;
314337
break;

‎src/mscp.c

+12-3
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
#include <print.h>
1818
#include <strerrno.h>
1919
#include <mscp.h>
20+
#include <bwlimit.h>
2021

2122
#include <openbsd-compat/openbsd-compat.h>
2223

@@ -56,6 +57,8 @@ struct mscp {
5657
#define chunk_pool_is_ready(m) ((m)->chunk_pool_ready)
5758
#define chunk_pool_set_ready(m, b) ((m)->chunk_pool_ready = b)
5859

60+
struct bwlimit bw; /* bandwidth limit mechanism */
61+
5962
struct mscp_thread scan; /* mscp_thread for mscp_scan_thread() */
6063
};
6164

@@ -281,6 +284,12 @@ struct mscp *mscp_init(struct mscp_opts *o, struct mscp_ssh_opts *s)
281284
pr_notice("usable cpu cores:%s", b);
282285
}
283286

287+
if (bwlimit_init(&m->bw, o->bitrate, 100) < 0) { /* 50ms window (hardcoded) */
288+
priv_set_errv("bwlimit_init: %s", strerrno());
289+
goto free_out;
290+
}
291+
pr_notice("bitrate limit: %lu bps", o->bitrate);
292+
284293
return m;
285294

286295
free_out:
@@ -522,8 +531,8 @@ int mscp_checkpoint_load(struct mscp *m, const char *pathname)
522531

523532
int mscp_checkpoint_save(struct mscp *m, const char *pathname)
524533
{
525-
return checkpoint_save(pathname, m->direction, m->ssh_opts->login_name,
526-
m->remote, m->path_pool, m->chunk_pool);
534+
return checkpoint_save(pathname, m->direction, m->ssh_opts->login_name, m->remote,
535+
m->path_pool, m->chunk_pool);
527536
}
528537

529538
static void *mscp_copy_thread(void *arg);
@@ -712,7 +721,7 @@ void *mscp_copy_thread(void *arg)
712721
}
713722

714723
if ((t->ret = copy_chunk(c, src_sftp, dst_sftp, m->opts->nr_ahead,
715-
m->opts->buf_sz, m->opts->preserve_ts,
724+
m->opts->buf_sz, m->opts->preserve_ts, &m->bw,
716725
&t->copied_bytes)) < 0)
717726
break;
718727
}

0 commit comments

Comments
 (0)