1// SPDX-License-Identifier: GPL-2.0-or-later
2/* IPVS:        Power of Twos Choice Scheduling module
3 *
4 * Authors:     Darby Payne <darby.payne@applovin.com>
5 */
6
7#define KMSG_COMPONENT "IPVS"
8#define pr_fmt(fmt) KMSG_COMPONENT ": " fmt
9
10#include <linux/kernel.h>
11#include <linux/module.h>
12#include <linux/random.h>
13
14#include <net/ip_vs.h>
15
16/*    Power of Twos Choice scheduling, algorithm originally described by
17 *    Michael Mitzenmacher.
18 *
19 *    Randomly picks two destinations and picks the one with the least
20 *    amount of connections
21 *
22 *    The algorithm calculates a few variables
23 *    - total_weight = sum of all weights
24 *    - rweight1 = random number between [0,total_weight]
25 *    - rweight2 = random number between [0,total_weight]
26 *
27 *    For each destination
28 *      decrement rweight1 and rweight2 by the destination weight
29 *      pick choice1 when rweight1 is <= 0
30 *      pick choice2 when rweight2 is <= 0
31 *
32 *    Return choice2 if choice2 has less connections than choice 1 normalized
33 *    by weight
34 *
35 * References
36 * ----------
37 *
38 * [Mitzenmacher 2016]
39 *    The Power of Two Random Choices: A Survey of Techniques and Results
40 *    Michael Mitzenmacher, Andrea W. Richa y, Ramesh Sitaraman
41 *    http://www.eecs.harvard.edu/~michaelm/NEWWORK/postscripts/twosurvey.pdf
42 *
43 */
44static struct ip_vs_dest *ip_vs_twos_schedule(struct ip_vs_service *svc,
45					      const struct sk_buff *skb,
46					      struct ip_vs_iphdr *iph)
47{
48	struct ip_vs_dest *dest, *choice1 = NULL, *choice2 = NULL;
49	int rweight1, rweight2, weight1 = -1, weight2 = -1, overhead1 = 0;
50	int overhead2, total_weight = 0, weight;
51
52	IP_VS_DBG(6, "%s(): Scheduling...\n", __func__);
53
54	/* Generate a random weight between [0,sum of all weights) */
55	list_for_each_entry_rcu(dest, &svc->destinations, n_list) {
56		if (!(dest->flags & IP_VS_DEST_F_OVERLOAD)) {
57			weight = atomic_read(&dest->weight);
58			if (weight > 0) {
59				total_weight += weight;
60				choice1 = dest;
61			}
62		}
63	}
64
65	if (!choice1) {
66		ip_vs_scheduler_err(svc, "no destination available");
67		return NULL;
68	}
69
70	/* Add 1 to total_weight so that the random weights are inclusive
71	 * from 0 to total_weight
72	 */
73	total_weight += 1;
74	rweight1 = get_random_u32_below(total_weight);
75	rweight2 = get_random_u32_below(total_weight);
76
77	/* Pick two weighted servers */
78	list_for_each_entry_rcu(dest, &svc->destinations, n_list) {
79		if (dest->flags & IP_VS_DEST_F_OVERLOAD)
80			continue;
81
82		weight = atomic_read(&dest->weight);
83		if (weight <= 0)
84			continue;
85
86		rweight1 -= weight;
87		rweight2 -= weight;
88
89		if (rweight1 <= 0 && weight1 == -1) {
90			choice1 = dest;
91			weight1 = weight;
92			overhead1 = ip_vs_dest_conn_overhead(dest);
93		}
94
95		if (rweight2 <= 0 && weight2 == -1) {
96			choice2 = dest;
97			weight2 = weight;
98			overhead2 = ip_vs_dest_conn_overhead(dest);
99		}
100
101		if (weight1 != -1 && weight2 != -1)
102			goto nextstage;
103	}
104
105nextstage:
106	if (choice2 && (weight2 * overhead1) > (weight1 * overhead2))
107		choice1 = choice2;
108
109	IP_VS_DBG_BUF(6, "twos: server %s:%u conns %d refcnt %d weight %d\n",
110		      IP_VS_DBG_ADDR(choice1->af, &choice1->addr),
111		      ntohs(choice1->port), atomic_read(&choice1->activeconns),
112		      refcount_read(&choice1->refcnt),
113		      atomic_read(&choice1->weight));
114
115	return choice1;
116}
117
118static struct ip_vs_scheduler ip_vs_twos_scheduler = {
119	.name = "twos",
120	.refcnt = ATOMIC_INIT(0),
121	.module = THIS_MODULE,
122	.n_list = LIST_HEAD_INIT(ip_vs_twos_scheduler.n_list),
123	.schedule = ip_vs_twos_schedule,
124};
125
126static int __init ip_vs_twos_init(void)
127{
128	return register_ip_vs_scheduler(&ip_vs_twos_scheduler);
129}
130
131static void __exit ip_vs_twos_cleanup(void)
132{
133	unregister_ip_vs_scheduler(&ip_vs_twos_scheduler);
134	synchronize_rcu();
135}
136
137module_init(ip_vs_twos_init);
138module_exit(ip_vs_twos_cleanup);
139MODULE_LICENSE("GPL");
140MODULE_DESCRIPTION("ipvs power of twos choice scheduler");
141