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

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)
}