trident/flow_controller.h at master · GeekerClub/trident · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
// Copyright (c) 2014 The Trident Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.
//
//

#ifndef _TRIDENT_FLOW_CONTROLLER_H_
#define _TRIDENT_FLOW_CONTROLLER_H_

#include <trident/common_internal.h>

namespace trident {


class FlowController
{
public:
    FlowController(bool read_no_limit, int read_quota,
                   bool write_no_limit, int write_quota)
        : _read_no_limit(read_no_limit), _read_quota(read_quota)
        , _write_no_limit(write_no_limit), _write_quota(write_quota) {}

    ~FlowController() {}

    // Reset read quota.
    // @param read_no_limit  if set as no limit.
    // @param quota  the new quota, only useful when read_no_limit is false.
    void reset_read_quota(bool read_no_limit, int quota)
    {
        _read_no_limit = read_no_limit;
        _read_quota = quota;
    }

    // Reset read quota.
    // @param quota  the new quota.
    void reset_read_quota(int quota)
    {
        _read_quota = quota;
    }

    // Reset write quota.
    // @param write_no_limit  if set as no limit.
    // @param quota  the new quota, only useful when write_no_limit is false.
    void reset_write_quota(bool write_no_limit, int quota)
    {
        _write_no_limit = write_no_limit;
        _write_quota = quota;
    }

    // Reset write quota.
    // @param quota  the new quota.
    void reset_write_quota(int quota)
    {
        _write_quota = quota;
    }

    // Check if has more read quota.
    bool has_read_quota() const
    {
        return _read_no_limit || _read_quota > 0;
    }

    // Check if has more write quota.
    bool has_write_quota() const
    {
        return _write_no_limit || _write_quota > 0;
    }

    // Acquire some read quota.
    // @param quota  the quota expect to acquire.
    // @return >0  if acquire succeed.
    // @return <=0 if acquire failed, and the return value can be used as sequence number
    //             to sort trigger order: closer to zero, earlier to trigger.
    int acquire_read_quota(int quota)
    {
        return _read_no_limit ? 1 : atomic_add_ret_old(&_read_quota, -quota);
    }

    // Acquire some write quota.
    // @param quota  the quota expect to acquire.
    // @return >0  if acquire succeed.
    // @return <=0 if acquire failed, and the return value can be used as sequence number
    //             to sort trigger order: closer to zero, earlier to trigger.
    int acquire_write_quota(int quota)
    {
        return _write_no_limit ? 1 : atomic_add_ret_old(&_write_quota, -quota);
    }

private:
    bool _read_no_limit;
    volatile int _read_quota;
    bool _write_no_limit;
    volatile int _write_quota;

    TRIDENT_DISALLOW_EVIL_CONSTRUCTORS(FlowController);
}; // class FlowController


} // namespace trident

#endif // _TRIDENT_FLOW_CONTROLLER_H_

/* vim: set ts=4 sw=4 sts=4 tw=100 */