Browse Source

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

Loyalsoldier 4 years ago
parent
commit
de71e63893
2 changed files with 20 additions and 29 deletions
  1. 8 13
      common/geodata/cache.go
  2. 12 16
      common/geodata/decode.go

+ 8 - 13
common/geodata/cache.go

@@ -2,7 +2,6 @@ package geodata
 
 import (
 	"io/ioutil"
-	"runtime"
 	"strings"
 
 	"google.golang.org/protobuf/proto"
@@ -33,7 +32,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
 	}
@@ -48,11 +47,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 {
@@ -62,13 +61,11 @@ 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:
@@ -100,7 +97,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
 	}
@@ -115,11 +112,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 {
@@ -129,13 +126,11 @@ 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:

+ 12 - 16
common/geodata/decode.go

@@ -11,7 +11,6 @@ package geodata
 
 import (
 	"io"
-	"runtime"
 	"strings"
 
 	"google.golang.org/protobuf/encoding/protowire"
@@ -23,14 +22,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)
@@ -44,19 +43,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++
@@ -68,7 +67,7 @@ Loop:
 			}
 			lenVarint, n := protowire.ConsumeVarint(tempContainer)
 			if n < 0 {
-				return nil, errInvalidGeodataVarintLength
+				return nil, ErrInvalidGeodataVarintLength
 			}
 			tempContainer = nil
 			if !isInner {
@@ -98,11 +97,8 @@ 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
 }
 
@@ -113,7 +109,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
 	}