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