1/*
2 * Copyright (c) Huawei Technologies Co., Ltd. 2023. All rights reserved.
3 * Licensed under the Apache License, Version 2.0 (the "License");
4 * you may not use this file except in compliance with the License.
5 * You may obtain a copy of the License at
6 *
7 * http://www.apache.org/licenses/LICENSE-2.0
8 *
9 * Unless required by applicable law or agreed to in writing, software
10 * distributed under the License is distributed on an "AS IS" BASIS,
11 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 * See the License for the specific language governing permissions and
13 * limitations under the License.
14 */
15
16#ifndef __SAMPLING_H__
17#define __SAMPLING_H__
18
19#include <stdint.h>
20#include <cstddef>
21#include <random>
22#include <pthread.h>
23
24class Sampling {
25public:
26    void InitSampling(uint64_t);
27    void Reset();
28    uint64_t GetSampleInterval() const
29    {
30        return sampleInterval_;
31    }
32    size_t StartSampling(size_t allocSize)
33    {
34        return (allocSize >= sampleInterval_) ? allocSize : (CalcSamplings(allocSize) * sampleInterval_) ;
35    }
36    int64_t CalcNextSampleInterval()
37    {
38        return static_cast<int64_t>(exponentialDist_(randomEngine_)) + 1;
39    }
40    Sampling()
41    {
42        pthread_spin_init(&spinlock_, PTHREAD_PROCESS_PRIVATE);
43    }
44    ~Sampling()
45    {
46        pthread_spin_destroy(&spinlock_);
47    }
48
49    // for test
50    void SetRandomEngineSeed(uint64_t cnt)
51    {
52        randomEngine_.seed(cnt);
53    }
54
55private:
56    class Spinlock {
57    public:
58        Spinlock(const Spinlock&) = delete;
59        Spinlock& operator=(const Spinlock&) = delete;
60        Spinlock()
61        {
62            pthread_spin_lock(&spinlock_);
63        }
64        ~Spinlock()
65        {
66            pthread_spin_unlock(&spinlock_);
67        }
68    };
69    size_t CalcSamplings(size_t);
70
71    uint64_t sampleInterval_ {0};
72    double sampleRate_ {0.0};
73    int64_t nextSampleInterval_ {0};
74    std::exponential_distribution<double> exponentialDist_;
75    static inline std::default_random_engine randomEngine_;
76    static inline pthread_spinlock_t spinlock_;
77};
78#endif  // __SAMPLING_H__