updated the load balancer

hexahigh · May 1, 2024 · d82a9f9 · d82a9f9
1 parent ab71fae
commit d82a9f9
Show file tree

Hide file tree

Showing 5 changed files with 113 additions and 90 deletions.
diff --git a/backend/main.go b/backend/main.go
@@ -24,6 +24,7 @@ import (
 	"strconv"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"syscall"
 	"time"
 
@@ -70,6 +71,11 @@ const dbFilenamesSeperator = "||!??|"
 var db *sql.DB
 var logger *log.Logger
 
+var (
+	uploadCount   int64
+	downloadCount int64
+)
+
 //go:embed LICENSE
 var licenseFS embed.FS
 
@@ -156,6 +162,7 @@ func main() {
 	http.HandleFunc("/ping", handlePing)
 	http.HandleFunc("/health", handleHealth)
 	http.HandleFunc("/u/", handleU)
+	http.HandleFunc("/load", handleLoad)
 
 	if !*disableUpload {
 		http.HandleFunc("/store", handleStore)
@@ -217,6 +224,9 @@ func handleExists(w http.ResponseWriter, r *http.Request) {
 }
 
 func handleStore(w http.ResponseWriter, r *http.Request) {
+	atomic.AddInt64(&uploadCount, 1)
+	defer atomic.AddInt64(&uploadCount, -1)
+
 	type StoreResponse struct {
 		SHA256 string `json:"sha256"`
 		SHA1   string `json:"sha1"`
@@ -418,6 +428,9 @@ func handleStore(w http.ResponseWriter, r *http.Request) {
 }
 
 func handleGet(w http.ResponseWriter, r *http.Request) {
+	atomic.AddInt64(&downloadCount, 1)
+	defer atomic.AddInt64(&downloadCount, -1)
+
 	enableCors(&w)
 	if r.Method == "OPTIONS" {
 		return
@@ -457,6 +470,9 @@ func handleGet(w http.ResponseWriter, r *http.Request) {
 }
 
 func handleGet2(w http.ResponseWriter, r *http.Request) {
+	atomic.AddInt64(&downloadCount, 1)
+	defer atomic.AddInt64(&downloadCount, -1)
+
 	enableCors(&w)
 	if r.Method == "OPTIONS" {
 		return
@@ -753,6 +769,22 @@ func handlePing(w http.ResponseWriter, r *http.Request) {
 	w.Write([]byte("pong"))
 }
 
+func handleLoad(w http.ResponseWriter, r *http.Request) {
+	enableCors(&w)
+	if r.Method == "OPTIONS" {
+		return
+	}
+
+	response := map[string]interface{}{
+		"uploads":   atomic.LoadInt64(&uploadCount),
+		"downloads": atomic.LoadInt64(&downloadCount),
+	}
+
+	w.WriteHeader(http.StatusOK)
+	w.Header().Set("Content-Type", "application/json")
+	json.NewEncoder(w).Encode(response)
+}
+
 func dbFixer() {
 	logLevelln(0, "Cleaning database")
 

diff --git a/docs/api/v2/openapi.json b/docs/api/v2/openapi.json
@@ -243,12 +243,6 @@
                       "percentageUsed": {
                         "type": "number"
                       },
-                      "compression": {
-                        "type": "boolean"
-                      },
-                      "compression_level": {
-                        "type": "integer"
-                      },
                       "version": {
                         "type": "string"
                       },

diff --git a/frontend/src/lib/conf.js b/frontend/src/lib/conf.js
@@ -1,7 +1,7 @@
 export const endpoint = 'https://pomf1.080609.xyz'; // Main endpoint url
 export const instanceName = 'YAPC'; // instance name, shown in tos
 export const domain = 'pomf.080609.xyz'; // Domain where the instance is hosted, used in the TOS
-export const email = '[email protected]'; // Contact email. PLEASE change this if you are hosting your own instance (I do not want to get a DMCA takedown)
+export const email = '[email protected]'; // Contact email. PLEASE change this if you are hosting your own instance
 export const endpointList = { // list of endpoints displayed in the footer
     1: {
         name: "Local (8080)",
@@ -15,4 +15,13 @@ export const endpointList = { // list of endpoints displayed in the footer
         name: "Zerotier",
         url: 'http://10.0.0.5:9066',
     },
-}
+}
+
+export const loadEndpoints = [ // Used in the load balancer
+    {
+        name: 'NO1',
+        url: 'https://pomf1.080609.xyz',
+        lat: '59.2083',
+        lon: '10.9484'
+    }
+]
diff --git a/frontend/src/routes/+page.svelte b/frontend/src/routes/+page.svelte
@@ -45,13 +45,10 @@
 		const data = await response.json();
 		totalFiles = data.totalFiles || 'unknown';
 		totalSize = prettyBytes(data.totalSize) || 'unknown';
-		compression = data.compression;
-		compressionLevel = data.compression_level || 'unknown';
 		server_version = data.version || 'unknown';
 		freeSpace = prettyBytes(data.availableSpace);
 		totalSpace = prettyBytes(data.totalSpace);
 		percentageUsed = data.percentageUsed ? parseFloat(data.percentageUsed).toFixed(2) : 'unknown';
-		averageSpeed = data.averageSpeed ? prettyBytes(data.averageSpeed) : 'unknown';
 	}
 
 	async function archive(url) {
@@ -297,14 +294,11 @@
 									</p>
 									<p class="text-base text-gray-500">Statistics:</p>
 									<p class="text-sm text-gray-500">Server version: {server_version}</p>
-									<p class="text-sm text-gray-500">Average server speed: {averageSpeed}</p>
 									<p class="text-sm text-gray-500">Total files: {totalFiles}</p>
 									<p class="text-sm text-gray-500">Total file size: {totalSize}</p>
 									<p class="text-sm text-gray-500">Free space: {freeSpace}</p>
 									<p class="text-sm text-gray-500">Total space: {totalSpace}</p>
 									<p class="text-sm text-gray-500">Percentage used: {percentageUsed}</p>
-									<p class="text-sm text-gray-500">Compression: {compression}</p>
-									<p class="text-sm text-gray-500">Compression level: {compressionLevel}</p>
 								</div>
 							</div>
 						</div>

diff --git a/frontend/src/routes/f2/+server.js b/frontend/src/routes/f2/+server.js
@@ -1,4 +1,4 @@
-import { endpoint } from '$lib/conf.js';
+import { endpoint, loadEndpoints as servers } from '$lib/conf.js';
 
 export async function GET({ url, request }) {
 	// Extract the HASH and the file extension from the url
@@ -8,90 +8,84 @@ export async function GET({ url, request }) {
 
 	let fileUrl;
 
-	const servers = [
-		{
-			name: 'NO1',
-			url: 'https://pomf1.080609.xyz',
-			lat: '59.2083',
-			lon: '10.9484'
-		}
-	];
-
-	const ip =
-		request.headers.get('x-forwarded-for') || request.headers.get('remote_addr') || 'unknown';
+	// Select the least loaded server
+	const leastLoadedServer = await selectLeastLoadedServer(servers);
 
-	if (ip !== 'unknown') {
-		// Use ip-api.com to get the latitude and longitude of the IP address
-		const response = await fetch(`http://ip-api.com/json/${ip}`);
-		const data = await response.json();
-		const clientLat = data.lat;
-		const clientLon = data.lon;
+	// Construct the URL to the file on the default server
+	fileUrl = `${leastLoadedServer}/get2/?h=${hash}&e=${ext}&f=${filename}`;
 
-		// Function to calculate the distance between two points using the Haversine formula
-		function calculateDistance(lat1, lon1, lat2, lon2) {
-			const R = 6371; // Radius of the earth in km
-			const dLat = deg2rad(lat2 - lat1);
-			const dLon = deg2rad(lon2 - lon1);
-			const a =
-				Math.sin(dLat / 2) * Math.sin(dLat / 2) +
-				Math.cos(deg2rad(lat1)) * Math.cos(deg2rad(lat2)) * Math.sin(dLon / 2) * Math.sin(dLon / 2);
-			const c = 2 * Math.atan2(Math.sqrt(a), Math.sqrt(1 - a));
-			const d = R * c; // Distance in km
-			return d;
+	// Create a 301 redirect response
+	return new Response(null, {
+		status: 301,
+		headers: {
+			Location: fileUrl
 		}
+	});
+}
 
-		function deg2rad(deg) {
-			return deg * (Math.PI / 180);
-		}
+async function selectLeastLoadedServer(servers, timeout = 2000) {
+	const loads = await Promise.all(
+		servers.map(async (server) => {
+			try {
+				const response = await fetchWithTimeout(`${server.url}/load`, {}, timeout);
+				if (!response.ok) {
+					throw new Error('Failed to fetch load');
+				}
+				const data = await response.json();
+				return { server: server.name, load: data.uploads + data.downloads };
+			} catch (error) {
+				console.error(`Failed to fetch load from ${server.name}:`, error);
+				return { server: server.name, load: Infinity }; // Return a high value to penalize this server
+			}
+		})
+	);
 
-		// Find the closest server
-		let closestServer = servers[0];
-		let shortestDistance = calculateDistance(
-			clientLat,
-			clientLon,
-			closestServer.lat,
-			closestServer.lon
-		);
+	// Check if all servers have failed (i.e., all loads are Infinity)
+	const allFailed = loads.every((load) => load.load === Infinity);
 
-		for (let server of servers) {
-			const distance = calculateDistance(clientLat, clientLon, server.lat, server.lon);
-			if (distance < shortestDistance) {
-				closestServer = server;
-				shortestDistance = distance;
-			}
-		}
+	if (allFailed) {
+		// Fallback to the default server
+		console.log('All servers failed to respond. Falling back to the default server.');
+		return endpoint;
+	}
 
-		// Function to send a request with a timeout
-		async function fetchWithTimeout(url, options, timeout = 2000) {
-			return Promise.race([
-				fetch(url, options),
-				new Promise((_, reject) => setTimeout(() => reject(new Error('Timeout')), timeout))
-			]);
-		}
+	let leastLoadedServer = servers[0].name;
+	let minLoad = loads[0].load;
 
-		try {
-			// Send a request to the closest server's health endpoint
-			await fetchWithTimeout(`${closestServer.url}/health`);
-		} catch (error) {
-			// If the request fails or times out, switch to another server
-			console.log('Failed to reach the closest server, switching to another server...');
-			// Implement logic to switch to another server here
-			// For simplicity, let's just select the next server in the list
-			closestServer = servers.find((server) => server.url !== closestServer.url) || servers[0];
+	for (let i = 1; i < loads.length; i++) {
+		if (loads[i].load < minLoad) {
+			leastLoadedServer = loads[i].server;
+			minLoad = loads[i].load;
 		}
-
-		// Construct the URL to the file on the selected server
-		fileUrl = `${closestServer.url}/get2/?h=${hash}&e=${ext}&f=${filename}`;
-	} else {
-		// Construct the URL to the file on the default server
-		fileUrl = `${endpoint}/get2/?h=${hash}&e=${ext}&f=${filename}`;
 	}
 
-	// Create a 301 redirect response
-	return new Response(null, {
-		status: 301,
-		headers: {
-			Location: fileUrl
-		}
-	});
+	// Find the server object for the least loaded server
+	const selectedServer = servers.find((server) => server.name === leastLoadedServer);
+
+	return selectedServer.url; // Return the URL of the least loaded server
+}
+
+// Function to calculate the distance between two points using the Haversine formula
+function calculateDistance(lat1, lon1, lat2, lon2) {
+	const R = 6371; // Radius of the earth in km
+	const dLat = deg2rad(lat2 - lat1);
+	const dLon = deg2rad(lon2 - lon1);
+	const a =
+		Math.sin(dLat / 2) * Math.sin(dLat / 2) +
+		Math.cos(deg2rad(lat1)) * Math.cos(deg2rad(lat2)) * Math.sin(dLon / 2) * Math.sin(dLon / 2);
+	const c = 2 * Math.atan2(Math.sqrt(a), Math.sqrt(1 - a));
+	const d = R * c; // Distance in km
+	return d;
+}
+
+function deg2rad(deg) {
+	return deg * (Math.PI / 180);
+}
+
+// Function to send a request with a timeout
+async function fetchWithTimeout(url, options, timeout = 2000) {
+	return Promise.race([
+		fetch(url, options),
+		new Promise((_, reject) => setTimeout(() => reject(new Error('Timeout')), timeout))
+	]);
 }