use go-cache

labring · Oct 11, 2023 · f397d27 · f397d27
1 parent e89f25c
commit f397d27
Show file tree

Hide file tree

Showing 4 changed files with 33 additions and 63 deletions.
diff --git a/controllers/admission/api/v1/icp.go b/controllers/admission/api/v1/icp.go
@@ -2,12 +2,13 @@ package v1
 
 import (
 	"encoding/json"
-	"golang.org/x/net/publicsuffix"
-	netv1 "k8s.io/api/networking/v1"
 	"net/http"
 	"net/url"
-	"sync"
 	"time"
+
+	"github.com/patrickmn/go-cache"
+	"golang.org/x/net/publicsuffix"
+	netv1 "k8s.io/api/networking/v1"
 )
 
 type IcpResponse struct {
@@ -23,102 +24,66 @@ type IcpResponse struct {
 	} `json:"result"`
 }
 
-type cacheEntry struct {
-	icpResponse *IcpResponse
-	cachedTime  time.Time
-	cacheTTL    time.Duration
-}
-
 type IcpValidator struct {
 	enabled  bool
 	endpoint string
 	key      string
 
-	cache sync.Map
+	cache *cache.Cache
 }
 
 func NewIcpValidator(icpEnabled bool, icpEndpoint string, icpKey string) *IcpValidator {
-	v := &IcpValidator{
+	return &IcpValidator{
 		enabled:  icpEnabled,
 		endpoint: icpEndpoint,
 		key:      icpKey,
-	}
-
-	// Start the cleanup routine in a goroutine
-	go v.cleanupRoutine()
-
-	return v
-}
-
-// cleanupRoutine periodically checks and deletes expired cache entries
-func (i *IcpValidator) cleanupRoutine() {
-	ticker := time.NewTicker(time.Hour) // checks every one hour
-	defer ticker.Stop()
-	for {
-		select {
-		case <-ticker.C:
-			i.cleanupExpiredCache()
-		}
+		cache:    cache.New(5*time.Minute, 3*time.Minute),
 	}
 }
 
-// cleanupExpiredCache checks all cache entries and deletes the expired ones
-func (i *IcpValidator) cleanupExpiredCache() {
-	i.cache.Range(func(key, value interface{}) bool {
-		entry := value.(cacheEntry)
-		if time.Since(entry.cachedTime) > entry.cacheTTL {
-			i.cache.Delete(key)
-		}
-		return true
-	})
-}
-
 func (i *IcpValidator) Query(rule *netv1.IngressRule) (*IcpResponse, error) {
 	domainName, err := publicsuffix.EffectiveTLDPlusOne(rule.Host)
 	if err != nil {
 		return nil, err
 	}
 
 	// Check if result is already cached
-	cache, found := i.cache.Load(domainName)
-	entry := cache.(cacheEntry)
-	// Use the cache entry if it's not expired
-	if found && time.Since(entry.cachedTime) <= entry.cacheTTL {
-		return entry.icpResponse, nil
-	} else if time.Since(entry.cachedTime) > entry.cacheTTL {
-		// Remove expired cache entry, and query again
-		i.cache.Delete(domainName)
+	cached, found := i.cache.Get(domainName)
+	if found {
+		return cached.(*IcpResponse), nil
 	}
 
 	// Query ICP
 	data := url.Values{}
 	data.Set("domainName", domainName)
 	data.Set("key", i.key)
+
 	resp, err := http.PostForm(i.endpoint, data)
 	if err != nil {
 		return nil, err
 	}
 	defer resp.Body.Close()
+
 	var response IcpResponse
 	if err := json.NewDecoder(resp.Body).Decode(&response); err != nil {
 		return nil, err
 	}
 
-	// Cache the result with the current timestamp, and generate a TTL by the response
-	i.cache.Store(domainName, cacheEntry{
-		icpResponse: &response,
-		cachedTime:  time.Now(),
-		cacheTTL:    genCacheTTL(&response),
-	})
+	// Cache the result with the current timestamp
+	i.cache.Set(
+		domainName,
+		&response,
+		genCacheTTL(&response),
+	)
 
 	return &response, nil
 }
 
 // genCacheTTL generates a cache TTL based on the response
 func genCacheTTL(rsp *IcpResponse) time.Duration {
-	// If the response is valid, and the site license is not empty, cache for 30 days
+	// If the response is valid, and the site license is not empty, cache for 24 hours
 	if rsp.ErrorCode == 0 && rsp.Result.SiteLicense != "" {
-		return 30 * 24 * time.Hour
+		return 24 * time.Hour
 	}
 	// Otherwise, cache for 5 minutes
 	return 5 * time.Minute

diff --git a/controllers/admission/api/v1/icp_test.go b/controllers/admission/api/v1/icp_test.go
@@ -7,12 +7,14 @@ import (
 
 func TestIcpValidator_Query(t *testing.T) {
 	icpValidator := NewIcpValidator(true, "http://v.juhe.cn/siteTools/app/NewDomain/query.php", "")
-	rule := &v1.IngressRule{
-		Host: "sealos.cn",
+	for i := 0; i <= 3; i++ {
+		rule := &v1.IngressRule{
+			Host: "sealos.cn",
+		}
+		icpResponse, err := icpValidator.Query(rule)
+		if err != nil {
+			t.Fatalf("Error querying ICP: %v", err)
+		}
+		t.Logf("ICP Response: %+v", icpResponse)
 	}
-	icpResponse, err := icpValidator.Query(rule)
-	if err != nil {
-		t.Fatalf("Error querying ICP: %v", err)
-	}
-	t.Logf("ICP Response: %+v", icpResponse)
 }
diff --git a/controllers/admission/go.mod b/controllers/admission/go.mod
@@ -14,6 +14,8 @@ require (
 	github.com/labring/sealos v0.0.0-00010101000000-000000000000
 	github.com/onsi/ginkgo/v2 v2.10.0
 	github.com/onsi/gomega v1.27.8
+	github.com/patrickmn/go-cache v2.1.0+incompatible
+	golang.org/x/net v0.10.0
 	k8s.io/api v0.27.4
 	k8s.io/apimachinery v0.27.4
 	k8s.io/client-go v0.27.4
@@ -70,7 +72,6 @@ require (
 	go.uber.org/multierr v1.11.0 // indirect
 	go.uber.org/zap v1.24.0 // indirect
 	golang.org/x/crypto v0.10.0 // indirect
-	golang.org/x/net v0.10.0 // indirect
 	golang.org/x/oauth2 v0.8.0 // indirect
 	golang.org/x/sys v0.9.0 // indirect
 	golang.org/x/term v0.9.0 // indirect

diff --git a/controllers/admission/go.sum b/controllers/admission/go.sum
@@ -148,6 +148,8 @@ github.com/onsi/ginkgo/v2 v2.10.0 h1:sfUl4qgLdvkChZrWCYndY2EAu9BRIw1YphNAzy1VNWs
 github.com/onsi/ginkgo/v2 v2.10.0/go.mod h1:UDQOh5wbQUlMnkLfVaIUMtQ1Vus92oM+P2JX1aulgcE=
 github.com/onsi/gomega v1.27.8 h1:gegWiwZjBsf2DgiSbf5hpokZ98JVDMcWkUiigk6/KXc=
 github.com/onsi/gomega v1.27.8/go.mod h1:2J8vzI/s+2shY9XHRApDkdgPo1TKT7P2u6fXeJKFnNQ=
+github.com/patrickmn/go-cache v2.1.0+incompatible h1:HRMgzkcYKYpi3C8ajMPV8OFXaaRUnok+kx1WdO15EQc=
+github.com/patrickmn/go-cache v2.1.0+incompatible/go.mod h1:3Qf8kWWT7OJRJbdiICTKqZju1ZixQ/KpMGzzAfe6+WQ=
 github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
 github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=