11. 用Rust手把手編寫一個wmproxy(代理,內網穿透等), 實現健康檢查

2023-10-12 09:01:09

11. 用Rust手把手編寫一個wmproxy(代理,內網穿透等), 實現健康檢查

專案 ++wmproxy++

gite: https://gitee.com/tickbh/wmproxy

github: https://github.com/tickbh/wmproxy

健康檢查的意義

健康檢查維持著系統的穩定執行, 極大的加速著服務的響應時間, 並保證伺服器不會把訊息包轉發到不能響應的伺服器上, 從而使系統快速穩定的運轉
在LINUX系統中,系統預設TCP建立連線超時時間為127秒。通常網路不可達或者網路連線被拒絕或者網路連線超時需要耗時的時長較長。此時會超成伺服器的響應時間變長很多,而且重複發起不可達的連線嘗試也在耗著大量的IO資源。
當健康檢查介入後,如果短時間內多次建立連線失敗,則暫時判定該地址不可達,狀態設定為不可達。如果此時接收到該地址的請求時直接返回錯誤。大大提高了響應的時間。
所以健康檢查是必不可少的存在。

如何實現

由於健康狀態需要呼叫的地方可能在任意處需要發起連線的地方,如果通過引數透傳也會涉及到多執行緒的資料共用,如Arc<Mutex<Data>>,取用的時候也是要通過鎖共用,且編碼的複雜度和理解成本急劇升高,所以此處健康檢查選用的是多執行緒共用的靜態處理變數。

Rust中的靜態變數

在Rust中,全域性變數可以分為兩種:

  • 編譯期初始化的全域性變數
  • 執行期初始化的全域性變數

編譯期初始化的全域性變數有:

const建立的常數,如 const MAX_ID:usize=usize::MAX/2;
static建立的靜態變數,如 static mut REQUEST_RECV:usize=0;

執行期初始化的全域性變數有lazy_static用於懶初始化。例如:

lazy_static! {
    static ref HEALTH_CHECK: RwLock<HealthCheck> = RwLock::new(HealthCheck::new(60, 3, 2));
}

此外還有

  • 實現你自己的執行時初始化:std::sync::Once + static mut T
  • 單執行緒執行時初始化的特殊情況:thread_local

我們此處維持一個HealthCheck的全域性變數,因為程式是多執行緒,用thread_local,無法共用其它執行緒的檢測,不條例預期,所以此處用讀寫鎖來保證全域性變數的正確性,讀寫鎖的特點是允許存在多個讀,但如果獲取寫必須保證唯一。

原始碼解析,暫時不做主動性的健康檢查

接下來我們看HealthCheck的定義

pub struct HealthCheck {
    /// 健康檢查的重置時間, 失敗超過該時間會重新檢查, 統一單位秒
    fail_timeout: usize,
    /// 最大失敗次數, 一定時間內超過該次數認為不可存取
    max_fails: usize,
    /// 最小上線次數, 到達這個次數被認為存活
    min_rises: usize,
    /// 記錄跟地址相關的資訊
    health_map: HashMap<SocketAddr, HealthRecord>,
}

/// 每個SocketAddr的記錄值
struct HealthRecord {
    /// 最後的記錄時間
    last_record: Instant,
    /// 失敗的恢復時間
    fail_timeout: Duration,
    /// 當前連續失敗的次數
    fall_times: usize,
    /// 當前連續成功的次數
    rise_times: usize,
    /// 當前的狀態
    failed: bool,
}

主要有最後記錄時間,失敗次數,成功次數,最大失敗懲罰時間等元素組成

我們通過函數is_fall_down判定是否是異常狀態,未檢查前預設為正常狀態,超出一定時間後,解除異常狀態。

/// 檢測狀態是否能連線
pub fn is_fall_down(addr: &SocketAddr) -> bool {
    // 唯讀,獲取讀鎖
    if let Ok(h) = HEALTH_CHECK.read() {
        if !h.health_map.contains_key(addr) {
            return false;
        }
        let value = h.health_map.get(&addr).unwrap();
        if Instant::now().duration_since(value.last_record) > value.fail_timeout {
            return false;
        }
        h.health_map[addr].failed
    } else {
        false
    }
}

如果連線TCP失敗則呼叫add_fall_down將該地址失敗連線次數+1,如果失敗次數達到最大失敗次數將狀態置為不可用。

/// 失敗時呼叫
pub fn add_fall_down(addr: SocketAddr) {
    // 需要寫入,獲取寫入鎖
    if let Ok(mut h) = HEALTH_CHECK.write() {
        if !h.health_map.contains_key(&addr) {
            let mut health = HealthRecord::new(h.fail_timeout);
            health.fall_times = 1;
            h.health_map.insert(addr, health);
        } else {
            let max_fails = h.max_fails;
            let value = h.health_map.get_mut(&addr).unwrap();
            // 超出最大的失敗時長,重新計算狀態
            if Instant::now().duration_since(value.last_record) > value.fail_timeout {
                value.clear_status();
            }
            value.last_record = Instant::now();
            value.fall_times += 1;
            value.rise_times = 0;

            if value.fall_times >= max_fails {
                value.failed = true;
            }
        }
    }
}

如果連線TCP成功則呼叫add_rise_up將該地址成功連線次數+1,如果成功次數達到最小次數將狀態置為不可用。

/// 成功時呼叫
pub fn add_rise_up(addr: SocketAddr) {
    // 需要寫入,獲取寫入鎖
    if let Ok(mut h) = HEALTH_CHECK.write() {
        if !h.health_map.contains_key(&addr) {
            let mut health = HealthRecord::new(h.fail_timeout);
            health.rise_times = 1;
            h.health_map.insert(addr, health);
        } else {
            let min_rises = h.min_rises;
            let value = h.health_map.get_mut(&addr).unwrap();
            // 超出最大的失敗時長,重新計算狀態
            if Instant::now().duration_since(value.last_record) > value.fail_timeout {
                value.clear_status();
            }
            value.last_record = Instant::now();
            value.rise_times += 1;
            value.fall_times = 0;

            if value.rise_times >= min_rises {
                value.failed = false;
            }
        }
    }
}

接下來我們將TcpStream::connect函數統一替換成HealthCheck::connect外部修改幾乎為0,可實現開啟健康檢查,後續還會有主動式的健康檢查。

pub async fn connect<A>(addr: &A) -> io::Result<TcpStream>
    where
        A: ToSocketAddrs,
    {
        let addrs = addr.to_socket_addrs()?;
        let mut last_err = None;

        for addr in addrs {
            // 健康檢查失敗,直接返回錯誤
            if Self::is_fall_down(&addr) {
                last_err = Some(io::Error::new(io::ErrorKind::Other, "health check falldown"));
            } else {
                match TcpStream::connect(&addr).await {
                    Ok(stream) => 
                    {
                        Self::add_rise_up(addr);
                        return Ok(stream)
                    },
                    Err(e) => {
                        Self::add_fall_down(addr);
                        last_err = Some(e)
                    },
                }
            }
        }

        Err(last_err.unwrap_or_else(|| {
            io::Error::new(
                io::ErrorKind::InvalidInput,
                "could not resolve to any address",
            )
        }))
    }

效果

在前三次請求的時候,將花費5秒左右才丟擲拒絕連結的錯誤

connect server Err(Os { code: 10061, kind: ConnectionRefused, message: "由於目標計算機積極拒絕,無
法連線。" })

可以發現三次之後,將會快速的丟擲錯誤,達成健康檢查的目標

connect server Err(Custom { kind: Other, error: "health check falldown" })

此時被動式的健康檢查已完成,後續按需要的話將按需看是否實現主動式的健康檢查。