chore: 并入 iotlcm-admin-svr 历史到 admin-svr/（原提交 49818c6）

admin-svr/server/internal/logic/memorydb/memorydb.go

@@ -0,0 +1,175 @@
+package memorydb
+
+import (
+	"context"
+	"errors"
+	"git.5iot.com/iotlcm-public/iotlcm-admin-svr/internal/cmd/inithttp"
+	"git.5iot.com/iotlcm-public/iotlcm-admin-svr/internal/library/mcache"
+	"git.5iot.com/iotlcm-public/iotlcm-admin-svr/internal/model"
+	"git.5iot.com/iotlcm-public/iotlcm-admin-svr/internal/service"
+	"sort"
+	"sync"
+	"time"
+
+	"github.com/gogf/gf/v2/database/gdb"
+	"github.com/gogf/gf/v2/errors/gerror"
+	"github.com/gogf/gf/v2/frame/g"
+)
+
+var (
+	db           gdb.DB
+	cacheUpdates sync.Mutex
+)
+
+type rebuildSourceContextKey struct{}
+
+func sourceContext(ctx context.Context) context.Context {
+	if source, ok := ctx.Value(rebuildSourceContextKey{}).(context.Context); ok {
+		return source
+	}
+	return ctx
+}
+
+type memDB interface {
+	Init(ctx context.Context) error
+	EventHandler(ctx context.Context) error
+	TableName(ctx context.Context) string
+	CreateTable(ctx context.Context) error
+	LoadData(ctx context.Context, pk int64) error
+	DeleteData(ctx context.Context, ids []int64) error
+}
+
+var memDBMap = make(map[string]memDB)
+
+func RegisterMemDB(name string, memDB memDB) {
+	memDBMap[name] = memDB
+}
+
+type sMemoryDB struct {
+	rebuildMu sync.Mutex
+}
+
+func NewMemoryDB() *sMemoryDB {
+	return &sMemoryDB{}
+}
+
+func init() {
+	service.RegisterMemoryDB(NewMemoryDB())
+	inithttp.RegisterHttpInitFunc("memorydb", func(ctx context.Context) error {
+		return service.MemoryDB().Init(ctx)
+	})
+}
+
+func (s *sMemoryDB) Init(ctx context.Context) error {
+	target, err := newMemoryDB(ctx)
+	if err != nil {
+		return err
+	}
+	if target == nil {
+		return nil
+	}
+	db = target
+	for _, name := range registeredMemoryDBNames() {
+		if err = memDBMap[name].Init(ctx); err != nil {
+			return err
+		}
+	}
+	g.Log().Debug(ctx, "sMemoryDB.Init end.", target.GetConfig())
+	return nil
+}
+
+// Rebuild replaces all cached tables in one SQLite transaction. The single
+// connection keeps readers from observing partially loaded tables, and any
+// failure rolls the entire rebuild back.
+func (s *sMemoryDB) Rebuild(ctx context.Context) error {
+	s.rebuildMu.Lock()
+	defer s.rebuildMu.Unlock()
+	cacheUpdates.Lock()
+	defer cacheUpdates.Unlock()
+
+	target := db
+	if target == nil {
+		return errors.New("memory database is unavailable")
+	}
+	if err := target.Transaction(ctx, func(txCtx context.Context, _ gdb.TX) error {
+		txCtx = context.WithValue(txCtx, rebuildSourceContextKey{}, ctx)
+		for _, name := range registeredMemoryDBNames() {
+			if err := memDBMap[name].CreateTable(txCtx); err != nil {
+				return gerror.Wrapf(err, "rebuild memory database table failed: %s", name)
+			}
+		}
+		return nil
+	}); err != nil {
+		return err
+	}
+	if err := mcache.Clear(ctx); err != nil {
+		return gerror.Wrap(err, "clear derived memory cache failed")
+	}
+	g.Log().Infof(ctx, "memory database rebuilt, tables: %d", len(memDBMap))
+	return nil
+}
+
+func newMemoryDB(ctx context.Context) (gdb.DB, error) {
+	link := g.Cfg().MustGet(ctx, "database.memorydb.link", "").String()
+	debug := g.Cfg().MustGet(ctx, "database.memorydb.debug", "false").Bool()
+	if link == "" {
+		g.Log().Warning(ctx, "memory database link is empty; queries will use the primary database")
+		return nil, nil
+	}
+	g.Log().Debug(ctx, "memory database link", link)
+	target, err := gdb.New(gdb.ConfigNode{
+		Link:             link,
+		MaxIdleConnCount: 1,
+		MaxOpenConnCount: 1,
+		MaxConnLifeTime:  24 * 3650 * time.Hour,
+		Debug:            debug,
+	})
+	if err != nil {
+		return nil, err
+	}
+	if target == nil {
+		return nil, errors.New("memory database is unavailable")
+	}
+	return target, nil
+}
+
+func registeredMemoryDBNames() []string {
+	names := make([]string, 0, len(memDBMap))
+	for name := range memDBMap {
+		names = append(names, name)
+	}
+	sort.Strings(names)
+	return names
+}
+
+func (s *sMemoryDB) EventHandler(ctx context.Context, chg *model.BroadcastDbChgCache) error {
+	cacheUpdates.Lock()
+	defer cacheUpdates.Unlock()
+	memDB, ok := memDBMap[chg.TableName]
+	if !ok {
+		return nil
+	}
+
+	if chg.Type == "add" || chg.Type == "edit" {
+		for _, pk := range chg.PK {
+			if err := memDB.LoadData(ctx, pk); err != nil {
+				return err
+			}
+		}
+	} else if chg.Type == "del" {
+		return memDB.DeleteData(ctx, chg.PK)
+	}
+	return nil
+}
+
+func (s *sMemoryDB) DB(ctx context.Context) gdb.DB {
+	link := g.Cfg().MustGet(ctx, "database.memorydb.link", "sqlite::@file(:memory:)").String()
+	if link == "" {
+		return g.DB()
+	}
+	return db
+}
+
+func (s *sMemoryDB) Ctx(ctx context.Context) gdb.DB {
+	return s.DB(ctx).Ctx(ctx)
+}