smbget: Adds a rate limiting option --limit-rate in KB/s

This patch implements a very simple rate limiter. It works by pausing the main
download loop whenever the bytes transferred are more than what we would get
with if it were transferred at the rate set by the user.
Please note that this may reduce the blocksize if the limit is too small.

Signed-off-by: Vasileios Porpodas <v.porpodas@gmail.com>
Reviewed-by: Jeremy Allison <jra@samba.org>
Reviewed-by: Andreas Schneider <asn@samba.org>
Autobuild-User(master): Jeremy Allison <jra@samba.org>
Autobuild-Date(master): Wed Nov  2 22:47:10 UTC 2022 on sn-devel-184

diff --git a/docs-xml/manpages/smbget.1.xml b/docs-xml/manpages/smbget.1.xml
index 64924ef9ca8cffd0f96ca0a020c0021f48862522..9d1db967f96fc347a9b9b27eb443d45d34d2f9ed 100644 (file)
--- a/docs-xml/manpages/smbget.1.xml
+++ b/docs-xml/manpages/smbget.1.xml
@@ -35,6 +35,7 @@
                <arg choice="opt">-O, --stdout</arg>
                <arg choice="opt">-u, --update</arg>
                <arg choice="opt">-e, --encrypt</arg>
+               <arg choice="opt">--limit-rate=INT</arg>
                <arg choice="opt">-?, --help</arg>
                <arg choice="opt">--usage</arg>
                <arg choice="req">smb://host/share/path/to/file</arg>
@@ -151,6 +152,11 @@
                <listitem><para>Enable SMB encryption.</para></listitem>
        </varlistentry>
 
+       <varlistentry>
+        <term>--limit-rate=INT</term>
+               <listitem><para>Limit download rate by this many KB/s.</para></listitem>
+       </varlistentry>
+
 </refsect1>
 
 <refsect1>

diff --git a/source3/script/tests/test_smbget.sh b/source3/script/tests/test_smbget.sh
index c1c99579a8c2e3a8ab618d40a23cc156b364438c..2322f3af942e12a120c639fc51d8e87a186a80f7 100755 (executable)
--- a/source3/script/tests/test_smbget.sh
+++ b/source3/script/tests/test_smbget.sh
@@ -269,6 +269,32 @@ test_msdfs_link()
        return 0
 }
 
+# Tests --limit-rate. Getting the testfile (128K in size) with --limit-rate 100
+# (that is 100KB/s) should take at least 1 sec to complete.
+test_limit_rate()
+{
+       clear_download_area
+       echo "$SMBGET -v -a --limit-rate 100 smb://$SERVER_IP/smbget/testfile"
+       time_begin=$(date +%s)
+       $SMBGET -v -a --limit-rate 100 smb://$SERVER_IP/smbget/testfile
+       if [ $? -ne 0 ]; then
+               echo 'ERROR: RC does not match, expected: 0'
+               return 1
+       fi
+       time_end=$(date +%s)
+       cmp --silent $WORKDIR/testfile ./testfile
+       if [ $? -ne 0 ]; then
+               echo 'ERROR: file content does not match'
+               return 1
+       fi
+       if [ $((time_end - time_begin)) -lt 1 ]; then
+               echo 'ERROR: It should take at least 1s to transfer 128KB with rate 100KB/s'
+               return 1
+       fi
+       return 0
+}
+
+
 create_test_data
 
 pushd $TMPDIR
@@ -306,6 +332,10 @@ testit "update" test_update ||
 
 testit "msdfs" test_msdfs_link ||
        failed=$((failed + 1))
+
+testit "limit rate" test_limit_rate ||
+       failed=$((failed + 1))
+
 clear_download_area
 
 popd # TMPDIR

diff --git a/source3/utils/smbget.c b/source3/utils/smbget.c
index 3e7c5687d83e9d6912b7794f4c5bf849d723a46e..5f3ac16b2042be57bca79b3904a72a3db03f3231 100644 (file)
--- a/source3/utils/smbget.c
+++ b/source3/utils/smbget.c
@@ -56,6 +56,7 @@ struct opt {
        bool send_stdout;
        bool update;
        int debuglevel;
+       unsigned limit_rate;
 };
 static struct opt opt = { .blocksize = SMB_DEFAULT_BLOCKSIZE };
 
@@ -355,6 +356,10 @@ static bool smb_download_file(const char *base, const char *name,
        off_t offset_download = 0, offset_check = 0, curpos = 0,
              start_offset = 0;
        struct stat localstat, remotestat;
+       clock_t start_of_bucket_ticks = 0;
+       size_t bytes_in_bucket = 0;
+       size_t bucket_size = 0;
+       clock_t ticks_to_fill_bucket = 0;
 
        snprintf(path, SMB_MAXPATHLEN-1, "%s%s%s", base,
                 (*base && *name && name[0] != '/' &&
@@ -576,6 +581,44 @@ static bool smb_download_file(const char *base, const char *name,
                offset_check = 0;
        }
 
+       /* We implement rate limiting by filling up a bucket with bytes and
+        * checking, once the bucket is filled, if it was filled too fast.
+        * If so, we sleep for some time to get an average transfer rate that
+        * equals to the one set by the user.
+        *
+        * The bucket size directly affects the traffic characteristics.
+        * The smaller the bucket the more frequent the pause/resume cycle.
+        * A large bucket can result in burst of high speed traffic and large
+        * pauses. A cycle of 100ms looks like a good value. This value (in
+        * ticks) is held in `ticks_to_fill_bucket`. The `bucket_size` is
+        * calculated as:
+        * `limit_rate * 1024 * / (CLOCKS_PER_SEC / ticks_to_fill_bucket)`
+        *
+        * After selecting the bucket size we also need to check the blocksize
+        * of the transfer, since this is the minimum unit of traffic that we
+        * can observe. Achieving a ~10% precision requires a blocksize with a
+        * maximum size of `bucket_size / 10`.
+        */
+       if (opt.limit_rate > 0) {
+               unsigned max_block_size;
+               /* This is the time that the bucket should take to fill. */
+               ticks_to_fill_bucket = 100 /*ms*/ * CLOCKS_PER_SEC / 1000;
+               /* This is the size of the bucket in bytes.
+                * If we fill the bucket too quickly we should pause */
+               bucket_size = opt.limit_rate * 1024 / (CLOCKS_PER_SEC / ticks_to_fill_bucket);
+               max_block_size = bucket_size / 10;
+               max_block_size = max_block_size > 0 ? max_block_size : 1;
+               if (opt.blocksize > max_block_size) {
+                       if (opt.blocksize != SMB_DEFAULT_BLOCKSIZE) {
+                               fprintf(stderr,
+                                       "Warning: Overriding block size to %d \
+                                        due to limit-rate", max_block_size);
+                       }
+                       opt.blocksize = max_block_size;
+               }
+               start_of_bucket_ticks = clock();
+       }
+
        readbuf = (char *)SMB_MALLOC(opt.blocksize);
        if (!readbuf) {
                fprintf(stderr, "Failed to allocate %zu bytes for read "
@@ -592,7 +635,30 @@ static bool smb_download_file(const char *base, const char *name,
                ssize_t bytesread;
                ssize_t byteswritten;
 
+               /* Rate limiting. This pauses the transfer to limit traffic. */
+               if (opt.limit_rate > 0) {
+                       if (bytes_in_bucket > bucket_size) {
+                               clock_t now_ticks = clock();
+                               clock_t diff_ticks = now_ticks
+                                                    - start_of_bucket_ticks;
+                               /* Check if the bucket filled up too fast. */
+                               if (diff_ticks < ticks_to_fill_bucket) {
+                                       /* Pause until `ticks_to_fill_bucket` */
+                                       double sleep_us
+                                        = (ticks_to_fill_bucket - diff_ticks)
+                                         * 1000000 / CLOCKS_PER_SEC;
+                                       usleep(sleep_us);
+                               }
+                               /* Reset the byte counter and the ticks. */
+                               bytes_in_bucket = 0;
+                               start_of_bucket_ticks = clock();
+                       }
+               }
+
                bytesread = smbc_read(remotehandle, readbuf, opt.blocksize);
+               if (opt.limit_rate > 0) {
+                       bytes_in_bucket += bytesread;
+               }
                if(bytesread < 0) {
                        fprintf(stderr,
                                "Can't read %zu bytes at offset %jd, file %s\n",
@@ -902,6 +968,13 @@ int main(int argc, char **argv)
                        .val        = 'f',
                        .descrip    = "Use specified rc file"
                },
+               {
+                       .longName   = "limit-rate",
+                       .argInfo    = POPT_ARG_INT,
+                       .arg        = &opt.limit_rate,
+                       .val        = 'l',
+                       .descrip    = "Limit download speed to this many KB/s"
+               },
 
                POPT_TABLEEND
        };