xref: /netbsd-src/external/bsd/zstd/dist/tests/rateLimiter.py (revision 3117ece4fc4a4ca4489ba793710b60b0d26bab6c)
1*3117ece4Schristos#!/usr/bin/env python3
2*3117ece4Schristos
3*3117ece4Schristos# ################################################################
4*3117ece4Schristos# Copyright (c) Meta Platforms, Inc. and affiliates.
5*3117ece4Schristos# All rights reserved.
6*3117ece4Schristos#
7*3117ece4Schristos# This source code is licensed under both the BSD-style license (found in the
8*3117ece4Schristos# LICENSE file in the root directory of this source tree) and the GPLv2 (found
9*3117ece4Schristos# in the COPYING file in the root directory of this source tree).
10*3117ece4Schristos# You may select, at your option, one of the above-listed licenses.
11*3117ece4Schristos# ##########################################################################
12*3117ece4Schristos
13*3117ece4Schristos# Rate limiter, replacement for pv
14*3117ece4Schristos# this rate limiter does not "catch up" after a blocking period
15*3117ece4Schristos# Limitations:
16*3117ece4Schristos# - only accepts limit speed in MB/s
17*3117ece4Schristos
18*3117ece4Schristosimport sys
19*3117ece4Schristosimport time
20*3117ece4Schristos
21*3117ece4SchristosMB = 1024 * 1024
22*3117ece4Schristosrate = float(sys.argv[1]) * MB
23*3117ece4Schristosstart = time.time()
24*3117ece4Schristostotal_read = 0
25*3117ece4Schristos
26*3117ece4Schristos# sys.stderr.close()  # remove error message, for Ctrl+C
27*3117ece4Schristos
28*3117ece4Schristostry:
29*3117ece4Schristos  buf = " "
30*3117ece4Schristos  while len(buf):
31*3117ece4Schristos    now = time.time()
32*3117ece4Schristos    to_read = max(int(rate * (now - start)), 1)
33*3117ece4Schristos    max_buf_size = 1 * MB
34*3117ece4Schristos    to_read = min(to_read, max_buf_size)
35*3117ece4Schristos    start = now
36*3117ece4Schristos
37*3117ece4Schristos    buf = sys.stdin.buffer.read(to_read)
38*3117ece4Schristos    sys.stdout.buffer.write(buf)
39*3117ece4Schristos
40*3117ece4Schristosexcept (KeyboardInterrupt, BrokenPipeError) as e:
41*3117ece4Schristos    pass
42