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