github.com/TeaOSLab/EdgeNode@v1.3.8/internal/iplibrary/ip_list_kv.go

// 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/idles"
	"github.com/TeaOSLab/EdgeNode/internal/utils/kvstore"
	"testing"
	"time"
)

type KVIPList struct {
	ipTable       *kvstore.Table[*pb.IPItem]
	versionsTable *kvstore.Table[int64]

	encoder *IPItemEncoder[*pb.IPItem]

	cleanTicker *time.Ticker

	isClosed bool

	offsetItemKey string
}

func NewKVIPList() (*KVIPList, error) {
	var db = &KVIPList{
		cleanTicker: time.NewTicker(24 * time.Hour),
		encoder:     &IPItemEncoder[*pb.IPItem]{},
	}
	err := db.init()
	return db, err
}

func (this *KVIPList) 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()
		})

		idles.RunTicker(this.cleanTicker, func() {
			if this.isClosed {
				return
			}
			deleteErr := this.DeleteExpiredItems()
			if deleteErr != nil {
				remotelogs.Error("IP_LIST_DB", "clean expired items failed: "+deleteErr.Error())
			}
		})
	})

	return nil
}

// Name 数据库名称代号
func (this *KVIPList) Name() string {
	return "kvstore"
}

// DeleteExpiredItems 删除过期的条目
func (this *KVIPList) 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 *KVIPList) 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 *KVIPList) 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 *KVIPList) 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 *KVIPList) UpdateMaxVersion(version int64) error {
	if this.isClosed {
		return nil
	}

	return this.versionsTable.Set("version", version)
}

func (this *KVIPList) 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 *KVIPList) Flush() error {
	return this.ipTable.DB().Store().Flush()
}

func (this *KVIPList) Close() error {
	this.isClosed = true
	return nil
}