From bbc892ad4b8e57bccfdb63191693c66db33960da Mon Sep 17 00:00:00 2001
From: GoEdgeLab <goedgecloud@gmail.com>
Date: Tue, 7 May 2024 17:55:28 +0800
Subject: [PATCH] =?UTF-8?q?=E4=B8=8A=E4=BC=A0IP=E5=A2=9E=E5=8A=A0=E5=A4=B1?=
 =?UTF-8?q?=E8=B4=A5=E9=87=8D=E8=AF=95?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 internal/waf/action_record_ip.go | 25 ++++++++++++++++++++-----
 1 file changed, 20 insertions(+), 5 deletions(-)

diff --git a/internal/waf/action_record_ip.go b/internal/waf/action_record_ip.go
index 2cad851..e7f01f5 100644
--- a/internal/waf/action_record_ip.go
+++ b/internal/waf/action_record_ip.go
@@ -9,6 +9,7 @@ import (
 	"github.com/TeaOSLab/EdgeNode/internal/goman"
 	"github.com/TeaOSLab/EdgeNode/internal/remotelogs"
 	"github.com/TeaOSLab/EdgeNode/internal/rpc"
+	memutils "github.com/TeaOSLab/EdgeNode/internal/utils/mem"
 	"github.com/TeaOSLab/EdgeNode/internal/waf/requests"
 	"github.com/iwind/TeaGo/types"
 	"net/http"
@@ -31,13 +32,22 @@ type recordIPTask struct {
 	sourceHTTPFirewallRuleSetId   int64
 }
 
-var recordIPTaskChan = make(chan *recordIPTask, 2048)
+var recordIPTaskChan = make(chan *recordIPTask, 512)
 
 func init() {
 	if !teaconst.IsMain {
 		return
 	}
 
+	var memGB = memutils.SystemMemoryGB()
+	if memGB > 16 {
+		recordIPTaskChan = make(chan *recordIPTask, 4<<10)
+	} else if memGB > 8 {
+		recordIPTaskChan = make(chan *recordIPTask, 2<<10)
+	} else if memGB > 4 {
+		recordIPTaskChan = make(chan *recordIPTask, 1<<10)
+	}
+
 	events.On(events.EventLoaded, func() {
 		goman.New(func() {
 			rpcClient, err := rpc.SharedRPC()
@@ -95,9 +105,15 @@ func init() {
 					for _, pbItem := range pbItemMap {
 						pbItems = append(pbItems, pbItem)
 					}
-					_, err = rpcClient.IPItemRPC.CreateIPItems(rpcClient.Context(), &pb.CreateIPItemsRequest{IpItems: pbItems})
-					if err != nil {
-						remotelogs.Error("WAF_RECORD_IP_ACTION", "create ip item failed: "+err.Error())
+
+					for i := 0; i < 5; /* max tries */ i++ {
+						_, err = rpcClient.IPItemRPC.CreateIPItems(rpcClient.Context(), &pb.CreateIPItemsRequest{IpItems: pbItems})
+						if err != nil {
+							remotelogs.Error("WAF_RECORD_IP_ACTION", "create ip item failed: "+err.Error())
+							time.Sleep(1 * time.Second)
+						} else {
+							break
+						}
 					}
 				} else {
 					time.Sleep(1 * time.Second)
@@ -137,7 +153,6 @@ func (this *RecordIPAction) WillChange() bool {
 func (this *RecordIPAction) Perform(waf *WAF, group *RuleGroup, set *RuleSet, request requests.Request, writer http.ResponseWriter) PerformResult {
 	var ipListId = this.IPListId
 
-
 	if ipListId <= 0 || firewallconfigs.IsGlobalListId(ipListId) {
 		// server or policy list ids
 		switch this.Type {