From 45f44c401acd31e0f5dcf8e7b4f2789cbc1be12a Mon Sep 17 00:00:00 2001
From: RPRX <63339210+rprx@users.noreply.github.com>
Date: Tue, 15 Dec 2020 20:27:18 +0800
Subject: [PATCH] Refactor: Optimize Memory Usage At Startup

https://github.com/XTLS/Xray-core/issues/68#issuecomment-745231528
---
 infra/conf/router.go | 124 ++++++++++++++++++++++++++++++++-----------
 main/run.go          |   7 +++
 2 files changed, 101 insertions(+), 30 deletions(-)

diff --git a/infra/conf/router.go b/infra/conf/router.go
index 236277d9..9bfc3aeb 100644
--- a/infra/conf/router.go
+++ b/infra/conf/router.go
@@ -2,6 +2,7 @@ package conf
 
 import (
 	"encoding/json"
+	"runtime"
 	"strconv"
 	"strings"
 
@@ -147,46 +148,109 @@ func ParseIP(s string) (*router.CIDR, error) {
 	}
 }
 
-func loadGeoIP(country string) ([]*router.CIDR, error) {
-	return loadIP("geoip.dat", country)
+func loadGeoIP(code string) ([]*router.CIDR, error) {
+	return loadIP("geoip.dat", code)
 }
 
-func loadIP(filename, country string) ([]*router.CIDR, error) {
-	geoipBytes, err := filesystem.ReadAsset(filename)
-	if err != nil {
-		return nil, newError("failed to open file: ", filename).Base(err)
-	}
-	var geoipList router.GeoIPList
-	if err := proto.Unmarshal(geoipBytes, &geoipList); err != nil {
-		return nil, err
-	}
+var (
+	FileCache = make(map[string][]byte)
+	IPCache   = make(map[string]*router.GeoIP)
+	SiteCache = make(map[string]*router.GeoSite)
+)
 
-	for _, geoip := range geoipList.Entry {
-		if geoip.CountryCode == country {
-			return geoip.Cidr, nil
+func loadFile(file string) ([]byte, error) {
+	if FileCache[file] == nil {
+		bs, err := filesystem.ReadAsset(file)
+		if err != nil {
+			return nil, newError("failed to open file: ", file).Base(err)
 		}
+		if len(bs) == 0 {
+			return nil, newError("empty file: ", file)
+		}
+		// Do not cache file, may save RAM when there
+		// are many files, but consume CPU each time.
+		return bs, nil
+		FileCache[file] = bs
 	}
-
-	return nil, newError("country not found in ", filename, ": ", country)
+	return FileCache[file], nil
 }
 
-func loadSite(filename, country string) ([]*router.Domain, error) {
-	geositeBytes, err := filesystem.ReadAsset(filename)
-	if err != nil {
-		return nil, newError("failed to open file: ", filename).Base(err)
-	}
-	var geositeList router.GeoSiteList
-	if err := proto.Unmarshal(geositeBytes, &geositeList); err != nil {
-		return nil, err
-	}
-
-	for _, site := range geositeList.Entry {
-		if site.CountryCode == country {
-			return site.Domain, nil
+func loadIP(file, code string) ([]*router.CIDR, error) {
+	index := file + ":" + code
+	if IPCache[index] == nil {
+		bs, err := loadFile(file)
+		if err != nil {
+			return nil, newError("failed to load file: ", file).Base(err)
 		}
+		bs = find(bs, []byte(code))
+		if bs == nil {
+			return nil, newError("code not found in ", file, ": ", code)
+		}
+		var geoip router.GeoIP
+		if err := proto.Unmarshal(bs, &geoip); err != nil {
+			return nil, newError("error unmarshal IP in ", file, ": ", code).Base(err)
+		}
+		defer runtime.GC()     // or debug.FreeOSMemory()
+		return geoip.Cidr, nil // do not cache geoip
+		IPCache[index] = &geoip
 	}
+	return IPCache[index].Cidr, nil
+}
 
-	return nil, newError("list not found in ", filename, ": ", country)
+func loadSite(file, code string) ([]*router.Domain, error) {
+	index := file + ":" + code
+	if SiteCache[index] == nil {
+		bs, err := loadFile(file)
+		if err != nil {
+			return nil, newError("failed to load file: ", file).Base(err)
+		}
+		bs = find(bs, []byte(code))
+		if bs == nil {
+			return nil, newError("list not found in ", file, ": ", code)
+		}
+		var geosite router.GeoSite
+		if err := proto.Unmarshal(bs, &geosite); err != nil {
+			return nil, newError("error unmarshal Site in ", file, ": ", code).Base(err)
+		}
+		defer runtime.GC()         // or debug.FreeOSMemory()
+		return geosite.Domain, nil // do not cache geosite
+		SiteCache[index] = &geosite
+	}
+	return SiteCache[index].Domain, nil
+}
+
+func find(data, code []byte) []byte {
+	codeL := len(code)
+	if codeL == 0 {
+		return nil
+	}
+	for {
+		dataL := len(data)
+		if dataL < 2 {
+			return nil
+		}
+		x, y := proto.DecodeVarint(data[1:])
+		if x == 0 && y == 0 {
+			return nil
+		}
+		headL, bodyL := 1+y, int(x)
+		dataL -= headL
+		if dataL < bodyL {
+			return nil
+		}
+		data = data[headL:]
+		if int(data[1]) == codeL {
+			for i := 0; i < codeL && data[2+i] == code[i]; i++ {
+				if i+1 == codeL {
+					return data[:bodyL]
+				}
+			}
+		}
+		if dataL == bodyL {
+			return nil
+		}
+		data = data[bodyL:]
+	}
 }
 
 type AttributeMatcher interface {
diff --git a/main/run.go b/main/run.go
index ebcc4b87..72bba5d1 100644
--- a/main/run.go
+++ b/main/run.go
@@ -9,12 +9,14 @@ import (
 	"path"
 	"path/filepath"
 	"runtime"
+	"runtime/debug"
 	"strings"
 	"syscall"
 
 	"github.com/xtls/xray-core/common/cmdarg"
 	"github.com/xtls/xray-core/common/platform"
 	"github.com/xtls/xray-core/core"
+	"github.com/xtls/xray-core/infra/conf"
 	"github.com/xtls/xray-core/main/commands/base"
 )
 
@@ -78,8 +80,13 @@ func executeRun(cmd *base.Command, args []string) {
 	}
 	defer server.Close()
 
+	conf.FileCache = nil
+	conf.IPCache = nil
+	conf.SiteCache = nil
+
 	// Explicitly triggering GC to remove garbage from config loading.
 	runtime.GC()
+	debug.FreeOSMemory()
 
 	{
 		osSignals := make(chan os.Signal, 1)