使用KV数据库来管理IP名单

internal/iplibrary/ip_list_db.go

@@ -1,305 +1,13 @@
-// Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+// Copyright 2024 GoEdge CDN goedge.cdn@gmail.com. All rights reserved. Official site: https://goedge.cn .
 
 package iplibrary
 
-import (
-	"github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb"
-	"github.com/TeaOSLab/EdgeNode/internal/events"
-	"github.com/TeaOSLab/EdgeNode/internal/goman"
-	"github.com/TeaOSLab/EdgeNode/internal/remotelogs"
-	"github.com/TeaOSLab/EdgeNode/internal/utils/dbs"
-	"github.com/iwind/TeaGo/Tea"
-	"os"
-	"path/filepath"
-	"time"
-)
+import "github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb"
 
-type IPListDB struct {
-	db *dbs.DB
-
-	itemTableName    string
-	versionTableName string
-
-	deleteExpiredItemsStmt   *dbs.Stmt
-	deleteItemStmt           *dbs.Stmt
-	insertItemStmt           *dbs.Stmt
-	selectItemsStmt          *dbs.Stmt
-	selectMaxItemVersionStmt *dbs.Stmt
-
-	selectVersionStmt *dbs.Stmt
-	updateVersionStmt *dbs.Stmt
-
-	cleanTicker *time.Ticker
-
-	dir string
-
-	isClosed bool
-}
-
-func NewIPListDB() (*IPListDB, error) {
-	var db = &IPListDB{
-		itemTableName:    "ipItems",
-		versionTableName: "versions",
-		dir:              filepath.Clean(Tea.Root + "/data"),
-		cleanTicker:      time.NewTicker(24 * time.Hour),
-	}
-	err := db.init()
-	return db, err
-}
-
-func (this *IPListDB) init() error {
-	// 检查目录是否存在
-	_, err := os.Stat(this.dir)
-	if err != nil {
-		err = os.MkdirAll(this.dir, 0777)
-		if err != nil {
-			return err
-		}
-		remotelogs.Println("IP_LIST_DB", "create data dir '"+this.dir+"'")
-	}
-
-	var path = this.dir + "/ip_list.db"
-
-	db, err := dbs.OpenWriter("file:" + path + "?cache=shared&mode=rwc&_journal_mode=WAL&_sync=" + dbs.SyncMode + "&_locking_mode=EXCLUSIVE")
-	if err != nil {
-		return err
-	}
-	db.SetMaxOpenConns(1)
-
-	//_, err = db.Exec("VACUUM")
-	//if err != nil {
-	//	return err
-	//}
-
-	this.db = db
-
-	// 恢复数据库
-	var recoverEnv, _ = os.LookupEnv("EdgeRecover")
-	if len(recoverEnv) > 0 {
-		for _, indexName := range []string{"ip_list_itemId", "ip_list_expiredAt"} {
-			_, _ = db.Exec(`REINDEX "` + indexName + `"`)
-		}
-	}
-
-	// 初始化数据库
-	_, err = db.Exec(`CREATE TABLE IF NOT EXISTS "` + this.itemTableName + `" (
-  "id" integer NOT NULL PRIMARY KEY AUTOINCREMENT,
-  "listId" integer DEFAULT 0,
-  "listType" varchar(32),
-  "isGlobal" integer(1) DEFAULT 0,
-  "type" varchar(16),
-  "itemId" integer DEFAULT 0,
-  "ipFrom" varchar(64) DEFAULT 0,
-  "ipTo" varchar(64) DEFAULT 0,
-  "expiredAt" integer DEFAULT 0,
-  "eventLevel" varchar(32),
-  "isDeleted" integer(1) DEFAULT 0,
-  "version" integer DEFAULT 0,
-  "nodeId" integer DEFAULT 0,
-  "serverId" integer DEFAULT 0
-);
-
-CREATE INDEX IF NOT EXISTS "ip_list_itemId"
-ON "` + this.itemTableName + `" (
-  "itemId" ASC
-);
-
-CREATE INDEX IF NOT EXISTS "ip_list_expiredAt"
-ON "` + this.itemTableName + `" (
-  "expiredAt" ASC
-);
-`)
-	if err != nil {
-		return err
-	}
-
-	_, err = db.Exec(`CREATE TABLE IF NOT EXISTS "` + this.versionTableName + `" (
-  "id" integer NOT NULL PRIMARY KEY AUTOINCREMENT,
-  "version" integer DEFAULT 0
-);
-`)
-	if err != nil {
-		return err
-	}
-
-	// 初始化SQL语句
-	this.deleteExpiredItemsStmt, err = this.db.Prepare(`DELETE FROM "` + this.itemTableName + `" WHERE  "expiredAt">0 AND "expiredAt"<?`)
-	if err != nil {
-		return err
-	}
-
-	this.deleteItemStmt, err = this.db.Prepare(`DELETE FROM "` + this.itemTableName + `" WHERE "itemId"=?`)
-	if err != nil {
-		return err
-	}
-
-	this.insertItemStmt, err = this.db.Prepare(`INSERT INTO "` + this.itemTableName + `" ("listId", "listType", "isGlobal", "type", "itemId", "ipFrom", "ipTo", "expiredAt", "eventLevel", "isDeleted", "version", "nodeId", "serverId") VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`)
-	if err != nil {
-		return err
-	}
-
-	this.selectItemsStmt, err = this.db.Prepare(`SELECT "listId", "listType", "isGlobal", "type", "itemId", "ipFrom", "ipTo", "expiredAt", "eventLevel", "isDeleted", "version", "nodeId", "serverId" FROM "` + this.itemTableName + `" WHERE isDeleted=0 ORDER BY "version" ASC, "itemId" ASC LIMIT ?, ?`)
-	if err != nil {
-		return err
-	}
-
-	this.selectMaxItemVersionStmt, err = this.db.Prepare(`SELECT "version" FROM "` + this.itemTableName + `" ORDER BY "id" DESC LIMIT 1`)
-	if err != nil {
-		return err
-	}
-
-	this.selectVersionStmt, err = this.db.Prepare(`SELECT "version" FROM "` + this.versionTableName + `" LIMIT 1`)
-	if err != nil {
-		return err
-	}
-
-	this.updateVersionStmt, err = this.db.Prepare(`REPLACE INTO "` + this.versionTableName + `" ("id", "version") VALUES (1, ?)`)
-	if err != nil {
-		return err
-	}
-
-	this.db = db
-
-	goman.New(func() {
-		events.OnClose(func() {
-			_ = this.Close()
-			this.cleanTicker.Stop()
-		})
-
-		for range this.cleanTicker.C {
-			err := this.DeleteExpiredItems()
-			if err != nil {
-				remotelogs.Error("IP_LIST_DB", "clean expired items failed: "+err.Error())
-			}
-		}
-	})
-
-	return nil
-}
-
-// DeleteExpiredItems 删除过期的条目
-func (this *IPListDB) DeleteExpiredItems() error {
-	if this.isClosed {
-		return nil
-	}
-
-	_, err := this.deleteExpiredItemsStmt.Exec(time.Now().Unix() - 7*86400)
-	return err
-}
-
-func (this *IPListDB) AddItem(item *pb.IPItem) error {
-	if this.isClosed {
-		return nil
-	}
-
-	_, err := this.deleteItemStmt.Exec(item.Id)
-	if err != nil {
-		return err
-	}
-
-	// 如果是删除，则不再创建新记录
-	if item.IsDeleted {
-		return this.UpdateMaxVersion(item.Version)
-	}
-
-	_, err = this.insertItemStmt.Exec(item.ListId, item.ListType, item.IsGlobal, item.Type, item.Id, item.IpFrom, item.IpTo, item.ExpiredAt, item.EventLevel, item.IsDeleted, item.Version, item.NodeId, item.ServerId)
-	if err != nil {
-		return err
-	}
-
-	return this.UpdateMaxVersion(item.Version)
-}
-
-func (this *IPListDB) ReadItems(offset int64, size int64) (items []*pb.IPItem, err error) {
-	if this.isClosed {
-		return
-	}
-
-	rows, err := this.selectItemsStmt.Query(offset, size)
-	if err != nil {
-		return nil, err
-	}
-	defer func() {
-		_ = rows.Close()
-	}()
-
-	for rows.Next() {
-		//  "listId", "listType", "isGlobal", "type", "itemId", "ipFrom", "ipTo", "expiredAt", "eventLevel", "isDeleted", "version", "nodeId", "serverId"
-		var pbItem = &pb.IPItem{}
-		err = rows.Scan(&pbItem.ListId, &pbItem.ListType, &pbItem.IsGlobal, &pbItem.Type, &pbItem.Id, &pbItem.IpFrom, &pbItem.IpTo, &pbItem.ExpiredAt, &pbItem.EventLevel, &pbItem.IsDeleted, &pbItem.Version, &pbItem.NodeId, &pbItem.ServerId)
-		if err != nil {
-			return nil, err
-		}
-		items = append(items, pbItem)
-	}
-	return
-}
-
-// ReadMaxVersion 读取当前最大版本号
-func (this *IPListDB) ReadMaxVersion() int64 {
-	if this.isClosed {
-		return 0
-	}
-
-	// from version table
-	{
-		var row = this.selectVersionStmt.QueryRow()
-		if row == nil {
-			return 0
-		}
-		var version int64
-		err := row.Scan(&version)
-		if err == nil {
-			return version
-		}
-	}
-
-	// from items table
-	{
-		var row = this.selectMaxItemVersionStmt.QueryRow()
-		if row == nil {
-			return 0
-		}
-		var version int64
-		err := row.Scan(&version)
-		if err != nil {
-			return 0
-		}
-
-		return version
-	}
-}
-
-// UpdateMaxVersion 修改版本号
-func (this *IPListDB) UpdateMaxVersion(version int64) error {
-	if this.isClosed {
-		return nil
-	}
-
-	_, err := this.updateVersionStmt.Exec(version)
-	return err
-}
-
-func (this *IPListDB) Close() error {
-	this.isClosed = true
-
-	if this.db != nil {
-		for _, stmt := range []*dbs.Stmt{
-			this.deleteExpiredItemsStmt,
-			this.deleteItemStmt,
-			this.insertItemStmt,
-			this.selectItemsStmt,
-			this.selectMaxItemVersionStmt, // ipItems table
-
-			this.selectVersionStmt, // versions table
-			this.updateVersionStmt,
-		} {
-			if stmt != nil {
-				_ = stmt.Close()
-			}
-		}
-
-		return this.db.Close()
-	}
-	return nil
+type IPListDB interface {
+	Name() string
+	DeleteExpiredItems() error
+	ReadMaxVersion() (int64, error)
+	ReadItems(offset int64, size int64) (items []*pb.IPItem, goNext bool, err error)
+	AddItem(item *pb.IPItem) error
 }

internal/iplibrary/ip_list_kv.go

@@ -0,0 +1,229 @@
+// Copyright 2024 GoEdge CDN goedge.cdn@gmail.com. All rights reserved. Official site: https://goedge.cn .
+
+package iplibrary
+
+import (
+	"encoding/binary"
+	"errors"
+	"github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb"
+	"github.com/TeaOSLab/EdgeNode/internal/events"
+	"github.com/TeaOSLab/EdgeNode/internal/goman"
+	"github.com/TeaOSLab/EdgeNode/internal/remotelogs"
+	"github.com/TeaOSLab/EdgeNode/internal/utils/fasttime"
+	"github.com/TeaOSLab/EdgeNode/internal/utils/kvstore"
+	"testing"
+	"time"
+)
+
+type IPListKV struct {
+	ipTable       *kvstore.Table[*pb.IPItem]
+	versionsTable *kvstore.Table[int64]
+
+	encoder *IPItemEncoder[*pb.IPItem]
+
+	cleanTicker *time.Ticker
+
+	isClosed bool
+
+	offsetItemKey string
+}
+
+func NewIPListKV() (*IPListKV, error) {
+	var db = &IPListKV{
+		cleanTicker: time.NewTicker(24 * time.Hour),
+		encoder:     &IPItemEncoder[*pb.IPItem]{},
+	}
+	err := db.init()
+	return db, err
+}
+
+func (this *IPListKV) init() error {
+	store, storeErr := kvstore.DefaultStore()
+	if storeErr != nil {
+		return storeErr
+	}
+	db, dbErr := store.NewDB("ip_list")
+	if dbErr != nil {
+		return dbErr
+	}
+
+	{
+		table, err := kvstore.NewTable[*pb.IPItem]("ip_items", this.encoder)
+		if err != nil {
+			return err
+		}
+		this.ipTable = table
+
+		err = table.AddFields("expiresAt")
+		if err != nil {
+			return err
+		}
+
+		db.AddTable(table)
+	}
+
+	{
+		table, err := kvstore.NewTable[int64]("versions", kvstore.NewIntValueEncoder[int64]())
+		if err != nil {
+			return err
+		}
+		this.versionsTable = table
+		db.AddTable(table)
+	}
+
+	goman.New(func() {
+		events.OnClose(func() {
+			_ = this.Close()
+			this.cleanTicker.Stop()
+		})
+
+		for range this.cleanTicker.C {
+			deleteErr := this.DeleteExpiredItems()
+			if deleteErr != nil {
+				remotelogs.Error("IP_LIST_DB", "clean expired items failed: "+deleteErr.Error())
+			}
+		}
+	})
+
+	return nil
+}
+
+// Name 数据库名称代号
+func (this *IPListKV) Name() string {
+	return "kvstore"
+}
+
+// DeleteExpiredItems 删除过期的条目
+func (this *IPListKV) DeleteExpiredItems() error {
+	if this.isClosed {
+		return nil
+	}
+
+	for {
+		var found bool
+		var currentTime = fasttime.Now().Unix()
+		err := this.ipTable.
+			Query().
+			FieldAsc("expiresAt").
+			ForUpdate().
+			Limit(1000).
+			FindAll(func(tx *kvstore.Tx[*pb.IPItem], item kvstore.Item[*pb.IPItem]) (goNext bool, err error) {
+				if !item.Value.IsDeleted && item.Value.ExpiredAt == 0 { // never expires
+					return kvstore.Skip()
+				}
+				if item.Value.ExpiredAt < currentTime-7*86400 /** keep for 7 days **/ {
+					err = tx.Delete(item.Key)
+					if err != nil {
+						return false, err
+					}
+					found = true
+					return true, nil
+				}
+
+				found = false
+				return false, nil
+			})
+		if err != nil {
+			return err
+		}
+		if !found {
+			break
+		}
+	}
+
+	return nil
+}
+
+func (this *IPListKV) AddItem(item *pb.IPItem) error {
+	if this.isClosed {
+		return nil
+	}
+
+	// 先删除
+	var key = this.encoder.EncodeKey(item)
+	err := this.ipTable.Delete(key)
+	if err != nil {
+		return err
+	}
+
+	// 如果是删除，则不再创建新记录
+	if item.IsDeleted {
+		return this.UpdateMaxVersion(item.Version)
+	}
+
+	err = this.ipTable.Set(key, item)
+	if err != nil {
+		return err
+	}
+
+	return this.UpdateMaxVersion(item.Version)
+}
+
+func (this *IPListKV) ReadItems(offset int64, size int64) (items []*pb.IPItem, goNextLoop bool, err error) {
+	if this.isClosed {
+		return
+	}
+
+	err = this.ipTable.
+		Query().
+		Offset(this.offsetItemKey).
+		Limit(int(size)).
+		FindAll(func(tx *kvstore.Tx[*pb.IPItem], item kvstore.Item[*pb.IPItem]) (goNext bool, err error) {
+			this.offsetItemKey = item.Key
+			goNextLoop = true
+
+			if !item.Value.IsDeleted {
+				items = append(items, item.Value)
+			}
+			return true, nil
+		})
+	return
+}
+
+// ReadMaxVersion 读取当前最大版本号
+func (this *IPListKV) ReadMaxVersion() (int64, error) {
+	if this.isClosed {
+		return 0, errors.New("database has been closed")
+	}
+
+	version, err := this.versionsTable.Get("version")
+	if err != nil {
+		if kvstore.IsNotFound(err) {
+			return 0, nil
+		}
+		return 0, err
+	}
+	return version, nil
+}
+
+// UpdateMaxVersion 修改版本号
+func (this *IPListKV) UpdateMaxVersion(version int64) error {
+	if this.isClosed {
+		return nil
+	}
+
+	return this.versionsTable.SetSync("version", version)
+}
+
+func (this *IPListKV) TestInspect(t *testing.T) error {
+	return this.ipTable.
+		Query().
+		FindAll(func(tx *kvstore.Tx[*pb.IPItem], item kvstore.Item[*pb.IPItem]) (goNext bool, err error) {
+			if len(item.Key) != 8 {
+				return false, errors.New("invalid key '" + item.Key + "'")
+			}
+
+			t.Log(binary.BigEndian.Uint64([]byte(item.Key)), "=>", item.Value)
+			return true, nil
+		})
+}
+
+// Flush to disk
+func (this *IPListKV) Flush() error {
+	return this.ipTable.DB().Store().Flush()
+}
+
+func (this *IPListKV) Close() error {
+	this.isClosed = true
+	return nil
+}

internal/iplibrary/ip_list_kv_objects.go

@@ -0,0 +1,55 @@
+// Copyright 2024 GoEdge CDN goedge.cdn@gmail.com. All rights reserved. Official site: https://goedge.cn .
+
+package iplibrary
+
+import (
+	"encoding/binary"
+	"errors"
+	"github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb"
+	"google.golang.org/protobuf/proto"
+	"math"
+)
+
+type IPItemEncoder[T interface{ *pb.IPItem }] struct {
+}
+
+func NewIPItemEncoder[T interface{ *pb.IPItem }]() *IPItemEncoder[T] {
+	return &IPItemEncoder[T]{}
+}
+
+func (this *IPItemEncoder[T]) Encode(value T) ([]byte, error) {
+	return proto.Marshal(any(value).(*pb.IPItem))
+}
+
+func (this *IPItemEncoder[T]) EncodeField(value T, fieldName string) ([]byte, error) {
+	switch fieldName {
+	case "expiresAt":
+		var expiresAt = any(value).(*pb.IPItem).ExpiredAt
+		if expiresAt < 0 || expiresAt > int64(math.MaxUint32) {
+			expiresAt = 0
+		}
+		var b = make([]byte, 4)
+		binary.BigEndian.PutUint32(b, uint32(expiresAt))
+		return b, nil
+	}
+
+	return nil, errors.New("field '" + fieldName + "' not found")
+}
+
+func (this *IPItemEncoder[T]) Decode(valueBytes []byte) (value T, err error) {
+	var item = &pb.IPItem{}
+	err = proto.Unmarshal(valueBytes, item)
+	value = item
+	return
+}
+
+// EncodeKey generate key for ip item
+func (this *IPItemEncoder[T]) EncodeKey(item *pb.IPItem) string {
+	var b = make([]byte, 8)
+	if item.Id < 0 {
+		item.Id = 0
+	}
+
+	binary.BigEndian.PutUint64(b, uint64(item.Id))
+	return string(b)
+}

internal/iplibrary/ip_list_kv_test.go

@@ -0,0 +1,221 @@
+// Copyright 2024 GoEdge CDN goedge.cdn@gmail.com. All rights reserved. Official site: https://goedge.cn .
+
+package iplibrary_test
+
+import (
+	"github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb"
+	"github.com/TeaOSLab/EdgeNode/internal/iplibrary"
+	"github.com/TeaOSLab/EdgeNode/internal/utils/fasttime"
+	"github.com/TeaOSLab/EdgeNode/internal/utils/testutils"
+	"github.com/TeaOSLab/EdgeNode/internal/zero"
+	"testing"
+	"time"
+)
+
+func TestIPListKV_AddItem(t *testing.T) {
+	kv, err := iplibrary.NewIPListKV()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	defer func() {
+		_ = kv.Flush()
+	}()
+
+	{
+		err = kv.AddItem(&pb.IPItem{
+			Id:        1,
+			IpFrom:    "192.168.1.101",
+			IpTo:      "",
+			Version:   1,
+			ExpiredAt: fasttime.NewFastTime().Unix() + 60,
+			ListId:    1,
+			IsDeleted: false,
+			ListType:  "white",
+		})
+		if err != nil {
+			t.Fatal(err)
+		}
+	}
+
+	{
+		err = kv.AddItem(&pb.IPItem{
+			Id:        2,
+			IpFrom:    "192.168.1.102",
+			IpTo:      "",
+			Version:   2,
+			ExpiredAt: fasttime.NewFastTime().Unix() + 60,
+			ListId:    1,
+			IsDeleted: false,
+			ListType:  "white",
+		})
+		if err != nil {
+			t.Fatal(err)
+		}
+	}
+
+	{
+		err = kv.AddItem(&pb.IPItem{
+			Id:        3,
+			IpFrom:    "192.168.1.103",
+			IpTo:      "",
+			Version:   3,
+			ExpiredAt: fasttime.NewFastTime().Unix() + 60,
+			ListId:    1,
+			IsDeleted: false,
+			ListType:  "white",
+		})
+		if err != nil {
+			t.Fatal(err)
+		}
+	}
+}
+
+func TestIPListKV_AddItems_Many(t *testing.T) {
+	kv, err := iplibrary.NewIPListKV()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	defer func() {
+		_ = kv.Flush()
+	}()
+
+	var count = 2
+	var from = 1
+	if testutils.IsSingleTesting() {
+		count = 2_000_000
+	}
+
+	var before = time.Now()
+	defer func() {
+		t.Logf("cost: %.2f s", time.Since(before).Seconds())
+	}()
+
+	for i := from; i <= from+count; i++ {
+		err = kv.AddItem(&pb.IPItem{
+			Id:        int64(i),
+			IpFrom:    testutils.RandIP(),
+			IpTo:      "",
+			Version:   int64(i),
+			ExpiredAt: fasttime.NewFastTime().Unix() + 86400,
+			ListId:    1,
+			IsDeleted: false,
+			ListType:  "white",
+		})
+		if err != nil {
+			t.Fatal(err)
+		}
+	}
+}
+
+func TestIPListKV_DeleteExpiredItems(t *testing.T) {
+	kv, err := iplibrary.NewIPListKV()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	defer func() {
+		_ = kv.Flush()
+	}()
+
+	err = kv.DeleteExpiredItems()
+	if err != nil {
+		t.Fatal(err)
+	}
+}
+
+func TestIPListKV_UpdateMaxVersion(t *testing.T) {
+	kv, err := iplibrary.NewIPListKV()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	defer func() {
+		_ = kv.Flush()
+	}()
+
+	err = kv.UpdateMaxVersion(101)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	maxVersion, err := kv.ReadMaxVersion()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	t.Log("version:", maxVersion)
+}
+
+func TestIPListKV_ReadMaxVersion(t *testing.T) {
+	kv, err := iplibrary.NewIPListKV()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	maxVersion, err := kv.ReadMaxVersion()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	t.Log("version:", maxVersion)
+}
+
+func TestIPListKV_ReadItems(t *testing.T) {
+	kv, err := iplibrary.NewIPListKV()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	for {
+		items, goNext, readErr := kv.ReadItems(0, 2)
+		if readErr != nil {
+			t.Fatal(readErr)
+		}
+		t.Log("====")
+		for _, item := range items {
+			t.Log(item.Id)
+		}
+
+		if !goNext {
+			break
+		}
+	}
+}
+
+func TestIPListKV_CountItems(t *testing.T) {
+	kv, err := iplibrary.NewIPListKV()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	var count int
+	var m = map[int64]zero.Zero{}
+	for {
+		items, goNext, readErr := kv.ReadItems(0, 1000)
+		if readErr != nil {
+			t.Fatal(readErr)
+		}
+		for _, item := range items {
+			count++
+			m[item.Id] = zero.Zero{}
+		}
+
+		if !goNext {
+			break
+		}
+	}
+	t.Log("count:", count, "len:", len(m))
+}
+
+func TestIPListKV_Inspect(t *testing.T) {
+	kv, err := iplibrary.NewIPListKV()
+	if err != nil {
+		t.Fatal(err)
+	}
+	err = kv.TestInspect(t)
+	if err != nil {
+		t.Fatal(err)
+	}
+}

internal/iplibrary/ip_list_sqlite.go

@@ -0,0 +1,312 @@
+// Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package iplibrary
+
+import (
+	"github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb"
+	"github.com/TeaOSLab/EdgeNode/internal/events"
+	"github.com/TeaOSLab/EdgeNode/internal/goman"
+	"github.com/TeaOSLab/EdgeNode/internal/remotelogs"
+	"github.com/TeaOSLab/EdgeNode/internal/utils/dbs"
+	"github.com/iwind/TeaGo/Tea"
+	"os"
+	"path/filepath"
+	"time"
+)
+
+type IPListSQLite struct {
+	db *dbs.DB
+
+	itemTableName    string
+	versionTableName string
+
+	deleteExpiredItemsStmt   *dbs.Stmt
+	deleteItemStmt           *dbs.Stmt
+	insertItemStmt           *dbs.Stmt
+	selectItemsStmt          *dbs.Stmt
+	selectMaxItemVersionStmt *dbs.Stmt
+
+	selectVersionStmt *dbs.Stmt
+	updateVersionStmt *dbs.Stmt
+
+	cleanTicker *time.Ticker
+
+	dir string
+
+	isClosed bool
+}
+
+func NewIPListSqlite() (*IPListSQLite, error) {
+	var db = &IPListSQLite{
+		itemTableName:    "ipItems",
+		versionTableName: "versions",
+		dir:              filepath.Clean(Tea.Root + "/data"),
+		cleanTicker:      time.NewTicker(24 * time.Hour),
+	}
+	err := db.init()
+	return db, err
+}
+
+func (this *IPListSQLite) init() error {
+	// 检查目录是否存在
+	_, err := os.Stat(this.dir)
+	if err != nil {
+		err = os.MkdirAll(this.dir, 0777)
+		if err != nil {
+			return err
+		}
+		remotelogs.Println("IP_LIST_DB", "create data dir '"+this.dir+"'")
+	}
+
+	var path = this.dir + "/ip_list.db"
+
+	db, err := dbs.OpenWriter("file:" + path + "?cache=shared&mode=rwc&_journal_mode=WAL&_sync=" + dbs.SyncMode + "&_locking_mode=EXCLUSIVE")
+	if err != nil {
+		return err
+	}
+	db.SetMaxOpenConns(1)
+
+	//_, err = db.Exec("VACUUM")
+	//if err != nil {
+	//	return err
+	//}
+
+	this.db = db
+
+	// 恢复数据库
+	var recoverEnv, _ = os.LookupEnv("EdgeRecover")
+	if len(recoverEnv) > 0 {
+		for _, indexName := range []string{"ip_list_itemId", "ip_list_expiredAt"} {
+			_, _ = db.Exec(`REINDEX "` + indexName + `"`)
+		}
+	}
+
+	// 初始化数据库
+	_, err = db.Exec(`CREATE TABLE IF NOT EXISTS "` + this.itemTableName + `" (
+  "id" integer NOT NULL PRIMARY KEY AUTOINCREMENT,
+  "listId" integer DEFAULT 0,
+  "listType" varchar(32),
+  "isGlobal" integer(1) DEFAULT 0,
+  "type" varchar(16),
+  "itemId" integer DEFAULT 0,
+  "ipFrom" varchar(64) DEFAULT 0,
+  "ipTo" varchar(64) DEFAULT 0,
+  "expiredAt" integer DEFAULT 0,
+  "eventLevel" varchar(32),
+  "isDeleted" integer(1) DEFAULT 0,
+  "version" integer DEFAULT 0,
+  "nodeId" integer DEFAULT 0,
+  "serverId" integer DEFAULT 0
+);
+
+CREATE INDEX IF NOT EXISTS "ip_list_itemId"
+ON "` + this.itemTableName + `" (
+  "itemId" ASC
+);
+
+CREATE INDEX IF NOT EXISTS "ip_list_expiredAt"
+ON "` + this.itemTableName + `" (
+  "expiredAt" ASC
+);
+`)
+	if err != nil {
+		return err
+	}
+
+	_, err = db.Exec(`CREATE TABLE IF NOT EXISTS "` + this.versionTableName + `" (
+  "id" integer NOT NULL PRIMARY KEY AUTOINCREMENT,
+  "version" integer DEFAULT 0
+);
+`)
+	if err != nil {
+		return err
+	}
+
+	// 初始化SQL语句
+	this.deleteExpiredItemsStmt, err = this.db.Prepare(`DELETE FROM "` + this.itemTableName + `" WHERE  "expiredAt">0 AND "expiredAt"<?`)
+	if err != nil {
+		return err
+	}
+
+	this.deleteItemStmt, err = this.db.Prepare(`DELETE FROM "` + this.itemTableName + `" WHERE "itemId"=?`)
+	if err != nil {
+		return err
+	}
+
+	this.insertItemStmt, err = this.db.Prepare(`INSERT INTO "` + this.itemTableName + `" ("listId", "listType", "isGlobal", "type", "itemId", "ipFrom", "ipTo", "expiredAt", "eventLevel", "isDeleted", "version", "nodeId", "serverId") VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`)
+	if err != nil {
+		return err
+	}
+
+	this.selectItemsStmt, err = this.db.Prepare(`SELECT "listId", "listType", "isGlobal", "type", "itemId", "ipFrom", "ipTo", "expiredAt", "eventLevel", "isDeleted", "version", "nodeId", "serverId" FROM "` + this.itemTableName + `" WHERE isDeleted=0 ORDER BY "version" ASC, "itemId" ASC LIMIT ?, ?`)
+	if err != nil {
+		return err
+	}
+
+	this.selectMaxItemVersionStmt, err = this.db.Prepare(`SELECT "version" FROM "` + this.itemTableName + `" ORDER BY "id" DESC LIMIT 1`)
+	if err != nil {
+		return err
+	}
+
+	this.selectVersionStmt, err = this.db.Prepare(`SELECT "version" FROM "` + this.versionTableName + `" LIMIT 1`)
+	if err != nil {
+		return err
+	}
+
+	this.updateVersionStmt, err = this.db.Prepare(`REPLACE INTO "` + this.versionTableName + `" ("id", "version") VALUES (1, ?)`)
+	if err != nil {
+		return err
+	}
+
+	this.db = db
+
+	goman.New(func() {
+		events.OnClose(func() {
+			_ = this.Close()
+			this.cleanTicker.Stop()
+		})
+
+		for range this.cleanTicker.C {
+			err := this.DeleteExpiredItems()
+			if err != nil {
+				remotelogs.Error("IP_LIST_DB", "clean expired items failed: "+err.Error())
+			}
+		}
+	})
+
+	return nil
+}
+
+// Name 数据库名称代号
+func (this *IPListSQLite) Name() string {
+	return "sqlite"
+}
+
+// DeleteExpiredItems 删除过期的条目
+func (this *IPListSQLite) DeleteExpiredItems() error {
+	if this.isClosed {
+		return nil
+	}
+
+	_, err := this.deleteExpiredItemsStmt.Exec(time.Now().Unix() - 7*86400)
+	return err
+}
+
+func (this *IPListSQLite) AddItem(item *pb.IPItem) error {
+	if this.isClosed {
+		return nil
+	}
+
+	_, err := this.deleteItemStmt.Exec(item.Id)
+	if err != nil {
+		return err
+	}
+
+	// 如果是删除，则不再创建新记录
+	if item.IsDeleted {
+		return this.UpdateMaxVersion(item.Version)
+	}
+
+	_, err = this.insertItemStmt.Exec(item.ListId, item.ListType, item.IsGlobal, item.Type, item.Id, item.IpFrom, item.IpTo, item.ExpiredAt, item.EventLevel, item.IsDeleted, item.Version, item.NodeId, item.ServerId)
+	if err != nil {
+		return err
+	}
+
+	return this.UpdateMaxVersion(item.Version)
+}
+
+func (this *IPListSQLite) ReadItems(offset int64, size int64) (items []*pb.IPItem, goNext bool, err error) {
+	if this.isClosed {
+		return
+	}
+
+	rows, err := this.selectItemsStmt.Query(offset, size)
+	if err != nil {
+		return nil, false, err
+	}
+	defer func() {
+		_ = rows.Close()
+	}()
+
+	for rows.Next() {
+		//  "listId", "listType", "isGlobal", "type", "itemId", "ipFrom", "ipTo", "expiredAt", "eventLevel", "isDeleted", "version", "nodeId", "serverId"
+		var pbItem = &pb.IPItem{}
+		err = rows.Scan(&pbItem.ListId, &pbItem.ListType, &pbItem.IsGlobal, &pbItem.Type, &pbItem.Id, &pbItem.IpFrom, &pbItem.IpTo, &pbItem.ExpiredAt, &pbItem.EventLevel, &pbItem.IsDeleted, &pbItem.Version, &pbItem.NodeId, &pbItem.ServerId)
+		if err != nil {
+			return nil, false, err
+		}
+		items = append(items, pbItem)
+	}
+
+	goNext = int64(len(items)) == size
+	return
+}
+
+// ReadMaxVersion 读取当前最大版本号
+func (this *IPListSQLite) ReadMaxVersion() (int64, error) {
+	if this.isClosed {
+		return 0, nil
+	}
+
+	// from version table
+	{
+		var row = this.selectVersionStmt.QueryRow()
+		if row == nil {
+			return 0, nil
+		}
+		var version int64
+		err := row.Scan(&version)
+		if err == nil {
+			return version, nil
+		}
+	}
+
+	// from items table
+	{
+		var row = this.selectMaxItemVersionStmt.QueryRow()
+		if row == nil {
+			return 0, nil
+		}
+		var version int64
+		err := row.Scan(&version)
+		if err != nil {
+			return 0, nil
+		}
+
+		return version, nil
+	}
+}
+
+// UpdateMaxVersion 修改版本号
+func (this *IPListSQLite) UpdateMaxVersion(version int64) error {
+	if this.isClosed {
+		return nil
+	}
+
+	_, err := this.updateVersionStmt.Exec(version)
+	return err
+}
+
+func (this *IPListSQLite) Close() error {
+	this.isClosed = true
+
+	if this.db != nil {
+		for _, stmt := range []*dbs.Stmt{
+			this.deleteExpiredItemsStmt,
+			this.deleteItemStmt,
+			this.insertItemStmt,
+			this.selectItemsStmt,
+			this.selectMaxItemVersionStmt, // ipItems table
+
+			this.selectVersionStmt, // versions table
+			this.updateVersionStmt,
+		} {
+			if stmt != nil {
+				_ = stmt.Close()
+			}
+		}
+
+		return this.db.Close()
+	}
+	return nil
+}

internal/iplibrary/ip_list_sqlite_test.go

@@ -12,7 +12,7 @@ import (
 )
 
 func TestIPListDB_AddItem(t *testing.T) {
-	db, err := iplibrary.NewIPListDB()
+	db, err := iplibrary.NewIPListSqlite()
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -59,7 +59,7 @@ func TestIPListDB_AddItem(t *testing.T) {
 }
 
 func TestIPListDB_ReadItems(t *testing.T) {
-	db, err := iplibrary.NewIPListDB()
+	db, err := iplibrary.NewIPListSqlite()
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -71,15 +71,16 @@ func TestIPListDB_ReadItems(t *testing.T) {
 		_ = db.Close()
 	}()
 
-	items, err := db.ReadItems(0, 2)
+	items, goNext, err := db.ReadItems(0, 2)
 	if err != nil {
 		t.Fatal(err)
 	}
+	t.Log("goNext:", goNext)
 	logs.PrintAsJSON(items, t)
 }
 
 func TestIPListDB_ReadMaxVersion(t *testing.T) {
-	db, err := iplibrary.NewIPListDB()
+	db, err := iplibrary.NewIPListSqlite()
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -90,7 +91,7 @@ func TestIPListDB_ReadMaxVersion(t *testing.T) {
 }
 
 func TestIPListDB_UpdateMaxVersion(t *testing.T) {
-	db, err := iplibrary.NewIPListDB()
+	db, err := iplibrary.NewIPListSqlite()
 	if err != nil {
 		t.Fatal(err)
 	}

internal/iplibrary/manager_ip_list.go

@@ -8,10 +8,13 @@ import (
 	"github.com/TeaOSLab/EdgeNode/internal/goman"
 	"github.com/TeaOSLab/EdgeNode/internal/remotelogs"
 	"github.com/TeaOSLab/EdgeNode/internal/rpc"
+	"github.com/TeaOSLab/EdgeNode/internal/trackers"
 	"github.com/TeaOSLab/EdgeNode/internal/utils"
 	"github.com/TeaOSLab/EdgeNode/internal/waf"
 	"github.com/TeaOSLab/EdgeNode/internal/zero"
 	"github.com/iwind/TeaGo/Tea"
+	"github.com/iwind/TeaGo/types"
+	"os"
 	"sync"
 	"time"
 )
@@ -45,7 +48,7 @@ func init() {
 type IPListManager struct {
 	ticker *time.Ticker
 
-	db *IPListDB
+	db IPListDB
 
 	lastVersion   int64
 	fetchPageSize int64
@@ -83,7 +86,7 @@ func (this *IPListManager) Start() {
 		case <-this.ticker.C:
 		case <-IPListUpdateNotify:
 		}
-		err := this.loop()
+		err = this.loop()
 		if err != nil {
 			countErrors++
 
@@ -110,7 +113,18 @@ func (this *IPListManager) Stop() {
 
 func (this *IPListManager) init() {
 	// 从数据库中当中读取数据
-	db, err := NewIPListDB()
+	// 检查sqlite文件是否存在，以便决定使用sqlite还是kv
+	var sqlitePath = Tea.Root + "/data/ip_list.db"
+	_, sqliteErr := os.Stat(sqlitePath)
+
+	var db IPListDB
+	var err error
+	if sqliteErr == nil {
+		db, err = NewIPListSqlite()
+	} else {
+		db, err = NewIPListKV()
+	}
+
 	if err != nil {
 		remotelogs.Error("IP_LIST_MANAGER", "create ip list local database failed: "+err.Error())
 	} else {
@@ -120,24 +134,30 @@ func (this *IPListManager) init() {
 		_ = db.DeleteExpiredItems()
 
 		// 本地数据库中最大版本号
-		this.lastVersion = db.ReadMaxVersion()
+		this.lastVersion, err = db.ReadMaxVersion()
+		if err != nil {
+			remotelogs.Error("IP_LIST_MANAGER", "find max version failed: "+err.Error())
+			this.lastVersion = 0
+		}
+		remotelogs.Println("IP_LIST_MANAGER", "starting from '"+db.Name()+"' version '"+types.String(this.lastVersion)+"' ...")
 
 		// 从本地数据库中加载
 		var offset int64 = 0
 		var size int64 = 2_000
+
+		var tr = trackers.Begin("IP_LIST_MANAGER:load")
+		defer tr.End()
+
 		for {
-			items, err := db.ReadItems(offset, size)
+			items, goNext, readErr := db.ReadItems(offset, size)
 			var l = len(items)
-			if err != nil {
-				remotelogs.Error("IP_LIST_MANAGER", "read ip list from local database failed: "+err.Error())
+			if readErr != nil {
+				remotelogs.Error("IP_LIST_MANAGER", "read ip list from local database failed: "+readErr.Error())
 			} else {
-				if l == 0 {
+				if !goNext {
 					break
 				}
 				this.processItems(items, false)
-				if int64(l) < size {
-					break
-				}
 			}
 			offset += int64(l)
 		}
@@ -310,9 +330,14 @@ func (this *IPListManager) processItems(items []*pb.IPItem, fromRemote bool) {
 
 // 调试IP信息
 func (this *IPListManager) debugItem(item *pb.IPItem) {
+	var ipRange = item.IpFrom
+	if len(item.IpTo) > 0 {
+		ipRange += " - " + item.IpTo
+	}
+
 	if item.IsDeleted {
-		remotelogs.Debug("IP_ITEM_DEBUG", "delete '"+item.IpFrom+"'")
+		remotelogs.Debug("IP_ITEM_DEBUG", "delete '"+ipRange+"'")
 	} else {
-		remotelogs.Debug("IP_ITEM_DEBUG", "add '"+item.IpFrom+"'")
+		remotelogs.Debug("IP_ITEM_DEBUG", "add '"+ipRange+"'")
 	}
 }