github.com/cilium/cilium@v1.16.2/bpf/lib/ratelimit.h (about)

     1  /* SPDX-License-Identifier: (GPL-2.0-only OR BSD-2-Clause) */
     2  /* Copyright Authors of Cilium */
     3  
     4  #pragma once
     5  
     6  #include "common.h"
     7  #include "bpf/helpers.h"
     8  
     9  struct ratelimit_key {
    10  	__u32 netdev_idx;
    11  };
    12  
    13  struct ratelimit_value {
    14  	__u64 last_topup;
    15  	__u64 tokens;
    16  };
    17  
    18  struct {
    19  	__uint(type, BPF_MAP_TYPE_LRU_HASH);
    20  	__type(key, struct ratelimit_key);
    21  	__type(value, struct ratelimit_value);
    22  	__uint(pinning, LIBBPF_PIN_BY_NAME);
    23  	__uint(max_entries, 1024);
    24  } RATELIMIT_MAP __section_maps_btf;
    25  
    26  struct ratelimit_settings {
    27  	/* A bucket will never have more than X amount of tokens, limits burt size */
    28  	__u64 bucket_size;
    29  	/* The amount of tokens added to a bucket for every topup */
    30  	__u64 tokens_per_topup;
    31  	/* The interval at which the topups happen */
    32  	__u64 topup_interval_ns;
    33  };
    34  
    35  static inline bool ratelimit_check_and_take(struct ratelimit_key *key,
    36  					    const struct ratelimit_settings *settings)
    37  {
    38  	struct ratelimit_value *value;
    39  	struct ratelimit_value new_value;
    40  	__u64 since_last_topup;
    41  	__u64 now;
    42  	__u64 interval;
    43  	__u64 remainder;
    44  	int ret;
    45  
    46  	now = ktime_get_ns();
    47  
    48  	/* Create a new bucket if we do not yet have one for the key */
    49  	value = map_lookup_elem(&RATELIMIT_MAP, key);
    50  	if (!value) {
    51  		new_value.last_topup = now;
    52  		new_value.tokens = settings->tokens_per_topup - 1;
    53  		ret = map_update_elem(&RATELIMIT_MAP, key, &new_value, BPF_ANY);
    54  		if (ret < 0)
    55  			return false;
    56  		return true;
    57  	}
    58  
    59  	/* Note, the updates below are racy, this causes a bit of inaccuracy but isn't fatal,
    60  	 * a more accurare implementation would use atomic operations to update the bucket
    61  	 * but this would be bad for performance.
    62  	 */
    63  
    64  	/* Topup the bucket if it has been at least more than 1 interval since we have done so */
    65  	since_last_topup = now - value->last_topup;
    66  	if (since_last_topup > settings->topup_interval_ns) {
    67  		interval = since_last_topup / settings->topup_interval_ns;
    68  		remainder = since_last_topup % settings->topup_interval_ns;
    69  		/* Add tokens of every missed interval */
    70  		value->tokens += interval * settings->tokens_per_topup;
    71  		value->last_topup = now - remainder;
    72  		/* Make sure to not overflow the bucket */
    73  		if (value->tokens > settings->bucket_size)
    74  			value->tokens = settings->bucket_size;
    75  	}
    76  
    77  	/* Take a token if there is at least one */
    78  	if (value->tokens > 0) {
    79  		value->tokens--;
    80  		return true;
    81  	}
    82  
    83  	return false;
    84  }