diff options
Diffstat (limited to 'src/zstd/tests/rateLimiter.py')
-rwxr-xr-x | src/zstd/tests/rateLimiter.py | 41 |
1 files changed, 41 insertions, 0 deletions
diff --git a/src/zstd/tests/rateLimiter.py b/src/zstd/tests/rateLimiter.py new file mode 100755 index 000000000..1068c4424 --- /dev/null +++ b/src/zstd/tests/rateLimiter.py @@ -0,0 +1,41 @@ +#!/usr/bin/env python3 + +# ################################################################ +# Copyright (c) 2018-2020, Facebook, Inc. +# All rights reserved. +# +# This source code is licensed under both the BSD-style license (found in the +# LICENSE file in the root directory of this source tree) and the GPLv2 (found +# in the COPYING file in the root directory of this source tree). +# You may select, at your option, one of the above-listed licenses. +# ########################################################################## + +# Rate limiter, replacement for pv +# this rate limiter does not "catch up" after a blocking period +# Limitations: +# - only accepts limit speed in MB/s + +import sys +import time + +MB = 1024 * 1024 +rate = float(sys.argv[1]) * MB +start = time.time() +total_read = 0 + +# sys.stderr.close() # remove error message, for Ctrl+C + +try: + buf = " " + while len(buf): + now = time.time() + to_read = max(int(rate * (now - start)), 1) + max_buf_size = 1 * MB + to_read = min(to_read, max_buf_size) + start = now + + buf = sys.stdin.buffer.read(to_read) + sys.stdout.buffer.write(buf) + +except (KeyboardInterrupt, BrokenPipeError) as e: + pass |