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

// Copyright 2019 TiKV Project Authors. Licensed under Apache-2.0.

use crate::grpc_sys::{self, grpc_resource_quota};
use std::ffi::CString;
use std::ptr;

/// ResourceQuota represents a bound on memory and thread usage by the gRPC.
/// NOTE: The management of threads created in grpc-core don't use ResourceQuota.
/// TODO: Manage the poller threads created in grpc-rs with this ResourceQuota later.
pub struct ResourceQuota {
    raw: *mut grpc_resource_quota,
}

impl ResourceQuota {
    /// Create a control block for resource quota. If a name is
    /// not declared for this control block, a name is automatically
    /// generated in grpc core.
    pub fn new(name: Option<&str>) -> ResourceQuota {
        match name {
            Some(name_str) => {
                let name_cstr = CString::new(name_str).unwrap();
                ResourceQuota {
                    raw: unsafe { grpc_sys::grpc_resource_quota_create(name_cstr.as_ptr() as _) },
                }
            }
            None => ResourceQuota {
                raw: unsafe { grpc_sys::grpc_resource_quota_create(ptr::null()) },
            },
        }
    }

    /// Resize this ResourceQuota to a new memory size.
    pub fn resize_memory(self, new_size: usize) -> ResourceQuota {
        unsafe { grpc_sys::grpc_resource_quota_resize(self.raw, new_size) };
        self
    }

    pub(crate) fn get_ptr(&self) -> *mut grpc_resource_quota {
        self.raw
    }
}

impl Drop for ResourceQuota {
    fn drop(&mut self) {
        unsafe {
            grpc_sys::grpc_resource_quota_unref(self.raw);
        }
    }
}