Revert "Refinement: geodata decoder removes unnecessary GC & exports methods for 3rd party (#965)"

This reverts commit de71e638

common/geodata/cache.go

@@ -2,6 +2,7 @@ package geodata
 
 import (
 	"io/ioutil"
+	"runtime"
 	"strings"
 
 	"google.golang.org/protobuf/proto"
@@ -32,7 +33,7 @@ func (g GeoIPCache) Set(key string, value *router.GeoIP) {
 
 func (g GeoIPCache) Unmarshal(filename, code string) (*router.GeoIP, error) {
 	asset := platform.GetAssetLocation(filename)
-	idx := strings.ToUpper(asset + "_" + code)
+	idx := strings.ToUpper(asset + "|" + code)
 	if g.Has(idx) {
 		return g.Get(idx), nil
 	}
@@ -47,11 +48,11 @@ func (g GeoIPCache) Unmarshal(filename, code string) (*router.GeoIP, error) {
 		g.Set(idx, &geoip)
 		return &geoip, nil
 
-	case ErrCodeNotFound:
+	case errCodeNotFound:
 		return nil, newError(code, " not found in ", filename)
 
-	case ErrFailedToReadBytes, ErrFailedToReadExpectedLenBytes,
-		ErrInvalidGeodataFile, ErrInvalidGeodataVarintLength:
+	case errFailedToReadBytes, errFailedToReadExpectedLenBytes,
+		errInvalidGeodataFile, errInvalidGeodataVarintLength:
 		newError("failed to decode geodata file: ", filename, ". Fallback to the original ReadFile method.").AtWarning().WriteToLog()
 		geoipBytes, err = ioutil.ReadFile(asset)
 		if err != nil {
@@ -61,11 +62,13 @@ func (g GeoIPCache) Unmarshal(filename, code string) (*router.GeoIP, error) {
 		if err := proto.Unmarshal(geoipBytes, &geoipList); err != nil {
 			return nil, err
 		}
+		runtime.GC()
 		for _, geoip := range geoipList.GetEntry() {
 			if strings.EqualFold(code, geoip.GetCountryCode()) {
 				g.Set(idx, geoip)
 				return geoip, nil
 			}
+			runtime.GC()
 		}
 
 	default:
@@ -97,7 +100,7 @@ func (g GeoSiteCache) Set(key string, value *router.GeoSite) {
 
 func (g GeoSiteCache) Unmarshal(filename, code string) (*router.GeoSite, error) {
 	asset := platform.GetAssetLocation(filename)
-	idx := strings.ToUpper(asset + "_" + code)
+	idx := strings.ToUpper(asset + "|" + code)
 	if g.Has(idx) {
 		return g.Get(idx), nil
 	}
@@ -112,11 +115,11 @@ func (g GeoSiteCache) Unmarshal(filename, code string) (*router.GeoSite, error)
 		g.Set(idx, &geosite)
 		return &geosite, nil
 
-	case ErrCodeNotFound:
+	case errCodeNotFound:
 		return nil, newError(code, " not found in ", filename)
 
-	case ErrFailedToReadBytes, ErrFailedToReadExpectedLenBytes,
-		ErrInvalidGeodataFile, ErrInvalidGeodataVarintLength:
+	case errFailedToReadBytes, errFailedToReadExpectedLenBytes,
+		errInvalidGeodataFile, errInvalidGeodataVarintLength:
 		newError("failed to decode geodata file: ", filename, ". Fallback to the original ReadFile method.").AtWarning().WriteToLog()
 		geositeBytes, err = ioutil.ReadFile(asset)
 		if err != nil {
@@ -126,11 +129,13 @@ func (g GeoSiteCache) Unmarshal(filename, code string) (*router.GeoSite, error)
 		if err := proto.Unmarshal(geositeBytes, &geositeList); err != nil {
 			return nil, err
 		}
+		runtime.GC()
 		for _, geosite := range geositeList.GetEntry() {
 			if strings.EqualFold(code, geosite.GetCountryCode()) {
 				g.Set(idx, geosite)
 				return geosite, nil
 			}
+			runtime.GC()
 		}
 
 	default:

common/geodata/decode.go

@@ -11,6 +11,7 @@ package geodata
 
 import (
 	"io"
+	"runtime"
 	"strings"
 
 	"google.golang.org/protobuf/encoding/protowire"
@@ -22,14 +23,14 @@ import (
 //go:generate go run github.com/v2fly/v2ray-core/v4/common/errors/errorgen
 
 var (
-	ErrFailedToReadBytes            = errors.New("failed to read bytes")
-	ErrFailedToReadExpectedLenBytes = errors.New("failed to read expected length of bytes")
-	ErrInvalidGeodataFile           = errors.New("invalid geodata file")
-	ErrInvalidGeodataVarintLength   = errors.New("invalid geodata varint length")
-	ErrCodeNotFound                 = errors.New("code not found")
+	errFailedToReadBytes            = errors.New("failed to read bytes")
+	errFailedToReadExpectedLenBytes = errors.New("failed to read expected length of bytes")
+	errInvalidGeodataFile           = errors.New("invalid geodata file")
+	errInvalidGeodataVarintLength   = errors.New("invalid geodata varint length")
+	errCodeNotFound                 = errors.New("code not found")
 )
 
-func EmitBytes(f io.ReadSeeker, code string) ([]byte, error) {
+func emitBytes(f io.ReadSeeker, code string) ([]byte, error) {
 	count := 1
 	isInner := false
 	tempContainer := make([]byte, 0, 5)
@@ -43,19 +44,19 @@ Loop:
 		container := make([]byte, advancedN)
 		bytesRead, err := f.Read(container)
 		if err == io.EOF {
-			return nil, ErrCodeNotFound
+			return nil, errCodeNotFound
 		}
 		if err != nil {
-			return nil, ErrFailedToReadBytes
+			return nil, errFailedToReadBytes
 		}
 		if bytesRead != len(container) {
-			return nil, ErrFailedToReadExpectedLenBytes
+			return nil, errFailedToReadExpectedLenBytes
 		}
 
 		switch count {
 		case 1, 3: // data type ((field_number << 3) | wire_type)
 			if container[0] != 10 { // byte `0A` equals to `10` in decimal
-				return nil, ErrInvalidGeodataFile
+				return nil, errInvalidGeodataFile
 			}
 			advancedN = 1
 			count++
@@ -67,7 +68,7 @@ Loop:
 			}
 			lenVarint, n := protowire.ConsumeVarint(tempContainer)
 			if n < 0 {
-				return nil, ErrInvalidGeodataVarintLength
+				return nil, errInvalidGeodataVarintLength
 			}
 			tempContainer = nil
 			if !isInner {
@@ -97,8 +98,11 @@ Loop:
 			result = container
 			break Loop
 		}
+
+		runtime.GC() // run GC every round to save memory
 	}
 
+	runtime.GC() // run GC at the end to save memory
 	return result, nil
 }
 
@@ -109,7 +113,7 @@ func Decode(filename, code string) ([]byte, error) {
 	}
 	defer f.Close()
 
-	geoBytes, err := EmitBytes(f, code)
+	geoBytes, err := emitBytes(f, code)
 	if err != nil {
 		return nil, err
 	}