vane 0.9.2

A flow-based reverse proxy with multi-layer routing and programmable pipelines.
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

/* src/layers/l4/balancer.rs */

use super::{
	health::{TARGET_HEALTH_REGISTRY, is_udp_target_healthy},
	model::{Forward, ResolvedTarget, Strategy},
	resolver,
};
use dashmap::DashMap;
use once_cell::sync::Lazy;
use rand::prelude::IndexedRandom;
use std::sync::atomic::{AtomicUsize, Ordering};

static SERIAL_COUNTERS: Lazy<DashMap<(u16, String), AtomicUsize>> = Lazy::new(DashMap::new);

/// Selects a TCP target by resolving, checking health, and applying a strategy.
pub async fn select_tcp_target(
	port: u16,
	rule_name: &str,
	forward_config: &Forward,
) -> Option<ResolvedTarget> {
	let resolved_targets = resolver::resolve_targets(&forward_config.targets).await;
	let available_targets: Vec<ResolvedTarget> = resolved_targets
		.into_iter()
		.filter(|t| TARGET_HEALTH_REGISTRY.get(t).is_some_and(|h| h.available))
		.collect();

	let chosen_pool = if !available_targets.is_empty() {
		available_targets
	} else {
		let resolved_fallbacks = resolver::resolve_targets(&forward_config.fallbacks).await;
		resolved_fallbacks
			.into_iter()
			.filter(|t| TARGET_HEALTH_REGISTRY.get(t).is_some_and(|h| h.available))
			.collect()
	};
	choose_from_pool(port, rule_name, &forward_config.strategy, chosen_pool)
}

/// Selects a UDP target by resolving, checking health, and applying a strategy.
pub async fn select_udp_target(
	port: u16,
	rule_name: &str,
	forward_config: &Forward,
) -> Option<ResolvedTarget> {
	let resolved_targets = resolver::resolve_targets(&forward_config.targets).await;
	let available_targets: Vec<ResolvedTarget> = resolved_targets
		.into_iter()
		.filter(is_udp_target_healthy)
		.collect();

	let chosen_pool = if !available_targets.is_empty() {
		available_targets
	} else {
		let resolved_fallbacks = resolver::resolve_targets(&forward_config.fallbacks).await;
		resolved_fallbacks
			.into_iter()
			.filter(is_udp_target_healthy)
			.collect()
	};
	choose_from_pool(port, rule_name, &forward_config.strategy, chosen_pool)
}

/// Chooses a target from a pool of resolved targets based on the configured strategy.
fn choose_from_pool(
	port: u16,
	rule_name: &str,
	strategy: &Strategy,
	pool: Vec<ResolvedTarget>,
) -> Option<ResolvedTarget> {
	if pool.is_empty() {
		return None;
	}
	match strategy {
		Strategy::Random => {
			let mut rng = rand::rng();
			pool.choose(&mut rng).cloned()
		}
		Strategy::Fastest => pool.into_iter().min_by_key(|t| {
			TARGET_HEALTH_REGISTRY
				.get(t)
				.map_or(std::time::Duration::MAX, |h| h.latency)
		}),
		Strategy::Serial => {
			let key = (port, rule_name.to_owned());
			let counter = SERIAL_COUNTERS.entry(key).or_default();
			let index = counter.fetch_add(1, Ordering::Relaxed) % pool.len();
			pool.get(index).cloned()
		}
	}
}