From b8b2490da6fab225c8d0d3650b83684212d4ce62 Mon Sep 17 00:00:00 2001
From: Wei Gao <gaowei66@huawei.com>
Date: Mon, 22 Feb 2021 09:31:16 +0800
Subject: [PATCH 3/3] qmp: implement qmp qos with leak_bucket.

1. Use leak bucket to limit handled qmp command number in unit of time.
Within 1 second, qmp command number will be limited in 100.

Signed-off-by: Wei Gao <gaowei66@huawei.com>
---
 docs/config_guidebook.md              |  4 ++
 machine_manager/src/qmp/mod.rs        | 28 +++++++++++-
 machine_manager/src/qmp/qmp_schema.rs |  5 +++
 machine_manager/src/socket.rs         | 62 ++++++++++++++++++++++-----
 4 files changed, 86 insertions(+), 13 deletions(-)

diff --git a/docs/config_guidebook.md b/docs/config_guidebook.md
index bc0f5d5..fe89671 100644
--- a/docs/config_guidebook.md
+++ b/docs/config_guidebook.md
@@ -566,6 +566,10 @@ When some events happen, connected client will receive QMP events.
 
 Now StratoVirt supports four events: `SHUTDOWN`, `STOP`, `RESUME`, `DEVICE_DELETED`.
 
+### 3.7 Flow control
+
+QMP use `leak bucket` to control QMP command flow. Now QMP server accept 100 commands per second.
+
 ## 4. Other Features
 
 ### 4.1 Daemonize
diff --git a/machine_manager/src/qmp/mod.rs b/machine_manager/src/qmp/mod.rs
index 0ced9fe..6b55fbc 100644
--- a/machine_manager/src/qmp/mod.rs
+++ b/machine_manager/src/qmp/mod.rs
@@ -44,11 +44,16 @@ use serde::{Deserialize, Serialize};
 use serde_json::Value;
 use vmm_sys_util::terminal::Terminal;
 
+use crate::event_loop::EventLoop;
 use crate::machine::MachineExternalInterface;
 use crate::socket::SocketRWHandler;
-use crate::{errors::Result, temp_cleaner::TempCleaner};
+use crate::{
+    errors::{Result, ResultExt},
+    temp_cleaner::TempCleaner,
+};
 use qmp_schema as schema;
 use schema::QmpCommand;
+use util::leak_bucket::LeakBucket;
 
 static mut QMP_CHANNEL: Option<Arc<QmpChannel>> = None;
 
@@ -323,12 +328,31 @@ pub fn create_timestamp() -> TimeStamp {
 ///
 /// * `stream_fd` - The input stream file description.
 /// * `controller` - The controller which execute actual qmp command.
+/// * `leak_bucket` - The LeakBucket flow controller for qmp command.
 ///
 /// # Errors
 ///
 /// This function will fail when json parser failed or socket file description broke.
-pub fn handle_qmp(stream_fd: RawFd, controller: &Arc<dyn MachineExternalInterface>) -> Result<()> {
+pub fn handle_qmp(
+    stream_fd: RawFd,
+    controller: &Arc<dyn MachineExternalInterface>,
+    leak_bucket: &mut LeakBucket,
+) -> Result<()> {
     let mut qmp_service = crate::socket::SocketHandler::new(stream_fd);
+
+    // If flow over `LEAK_BUCKET_LIMIT` per seconds, discard the request and return
+    // a `OperationThrottled` error.
+    if leak_bucket.throttled(EventLoop::get_ctx(None).unwrap()) {
+        qmp_service.discard()?;
+        let err_resp = schema::QmpErrorClass::OperationThrottled(crate::socket::LEAK_BUCKET_LIMIT);
+        qmp_service
+            .send_str(&serde_json::to_string(&Response::create_error_response(
+                err_resp, None,
+            ))?)
+            .chain_err(|| "Failed to send message to qmp client.")?;
+        return Ok(());
+    }
+
     match qmp_service.decode_line() {
         (Ok(None), _) => Ok(()),
         (Ok(buffer), if_fd) => {
diff --git a/machine_manager/src/qmp/qmp_schema.rs b/machine_manager/src/qmp/qmp_schema.rs
index 2b54642..5e4f408 100644
--- a/machine_manager/src/qmp/qmp_schema.rs
+++ b/machine_manager/src/qmp/qmp_schema.rs
@@ -31,6 +31,8 @@ pub enum QmpErrorClass {
     DeviceNotFound(String),
     #[serde(rename = "KVMMissingCap")]
     KVMMissingCap(String),
+    #[serde(rename = "OperationThrottled")]
+    OperationThrottled(u64),
 }
 
 impl QmpErrorClass {
@@ -41,6 +43,9 @@ impl QmpErrorClass {
             QmpErrorClass::DeviceNotActive(s) => s.to_string(),
             QmpErrorClass::DeviceNotFound(s) => s.to_string(),
             QmpErrorClass::KVMMissingCap(s) => s.to_string(),
+            QmpErrorClass::OperationThrottled(nr) => {
+                format!("More than {} requests received during 1 second", nr)
+            }
         }
     }
 }
diff --git a/machine_manager/src/socket.rs b/machine_manager/src/socket.rs
index 0d7a788..37dd176 100644
--- a/machine_manager/src/socket.rs
+++ b/machine_manager/src/socket.rs
@@ -16,7 +16,8 @@ use std::os::unix::io::{AsRawFd, RawFd};
 use std::os::unix::net::{UnixListener, UnixStream};
 use std::sync::{Arc, Mutex, RwLock};
 
-use util::loop_context::{EventNotifier, EventNotifierHelper, NotifierOperation};
+use util::leak_bucket::LeakBucket;
+use util::loop_context::{read_fd, EventNotifier, EventNotifierHelper, NotifierOperation};
 use vmm_sys_util::epoll::EventSet;
 
 use super::errors::Result;
@@ -27,6 +28,7 @@ use crate::{
 };
 
 const MAX_SOCKET_MSG_LENGTH: usize = 8192;
+pub(crate) const LEAK_BUCKET_LIMIT: u64 = 100;
 
 /// The wrapper over Unix socket and socket handler.
 ///
@@ -181,6 +183,10 @@ impl Socket {
         let mut notifiers = Vec::new();
         self.accept();
 
+        let leak_bucket = Arc::new(Mutex::new(LeakBucket::new(LEAK_BUCKET_LIMIT)));
+        let shared_leak_bucket = leak_bucket.clone();
+        let leak_bucket_fd = leak_bucket.lock().unwrap().as_raw_fd();
+
         let mut handlers = Vec::new();
         let handler: Box<dyn Fn(EventSet, RawFd) -> Option<Vec<EventNotifier>>> =
             Box::new(move |event, _| {
@@ -189,7 +195,11 @@ impl Socket {
                     let stream_fd = socket_mutexed.get_stream_fd();
 
                     let performer = &socket_mutexed.performer.as_ref().unwrap();
-                    if let Err(e) = crate::qmp::handle_qmp(stream_fd, performer) {
+                    if let Err(e) = crate::qmp::handle_qmp(
+                        stream_fd,
+                        performer,
+                        &mut shared_leak_bucket.lock().unwrap(),
+                    ) {
                         error!("{}", e);
                     }
                 }
@@ -200,28 +210,50 @@ impl Socket {
 
                     QmpChannel::unbind();
 
-                    Some(vec![EventNotifier::new(
-                        NotifierOperation::Delete,
-                        stream_fd,
-                        Some(listener_fd),
-                        EventSet::IN | EventSet::HANG_UP,
-                        Vec::new(),
-                    )])
+                    Some(vec![
+                        EventNotifier::new(
+                            NotifierOperation::Delete,
+                            stream_fd,
+                            Some(listener_fd),
+                            EventSet::IN | EventSet::HANG_UP,
+                            Vec::new(),
+                        ),
+                        EventNotifier::new(
+                            NotifierOperation::Delete,
+                            leak_bucket_fd,
+                            None,
+                            EventSet::IN,
+                            Vec::new(),
+                        ),
+                    ])
                 } else {
                     None
                 }
             });
         handlers.push(Arc::new(Mutex::new(handler)));
 
-        let notifier = EventNotifier::new(
+        let qmp_notifier = EventNotifier::new(
             NotifierOperation::AddShared,
             self.get_stream_fd(),
             Some(self.get_listener_fd()),
             EventSet::IN | EventSet::HANG_UP,
             handlers,
         );
+        notifiers.push(qmp_notifier);
+
+        let leak_bucket_notifier = EventNotifier::new(
+            NotifierOperation::AddShared,
+            leak_bucket_fd,
+            None,
+            EventSet::IN,
+            vec![Arc::new(Mutex::new(Box::new(move |_, fd| {
+                read_fd(fd);
+                leak_bucket.lock().unwrap().clear_timer();
+                None
+            })))],
+        );
+        notifiers.push(leak_bucket_notifier);
 
-        notifiers.push(notifier);
         notifiers
     }
 }
@@ -594,6 +626,14 @@ impl SocketHandler {
         }
     }
 
+    /// Discard message from `socket_fd`.
+    pub fn discard(&mut self) -> Result<()> {
+        self.stream.read_fd()?;
+        self.stream.clear();
+        self.buffer.clear();
+        Ok(())
+    }
+
     /// Send String to `socket_fd`.
     ///
     /// # Arguments
-- 
2.27.0

