Procházet zdrojové kódy

:art: 改进块树数据存取 https://github.com/siyuan-note/siyuan/issues/7168

Liang Ding před 2 roky
rodič
revize
fef57e49ca
3 změnil soubory, kde provedl 335 přidání a 185 odebrání
  1. 333 183
      kernel/treenode/blocktree.go
  2. 1 1
      kernel/util/working.go
  3. 1 1
      kernel/util/working_mobile.go

+ 333 - 183
kernel/treenode/blocktree.go

@@ -19,6 +19,7 @@ package treenode
 import (
 	"io"
 	"os"
+	"path/filepath"
 	"runtime"
 	"strings"
 	"sync"
@@ -28,14 +29,19 @@ import (
 	"github.com/88250/lute/ast"
 	"github.com/88250/lute/parse"
 	"github.com/dustin/go-humanize"
+	util2 "github.com/siyuan-note/dejavu/util"
 	"github.com/siyuan-note/logging"
 	"github.com/siyuan-note/siyuan/kernel/util"
 	"github.com/vmihailenco/msgpack/v5"
 )
 
-var blockTrees = map[string]*BlockTree{}
-var blockTreesLock = sync.Mutex{}
-var blockTreesChanged = time.Time{}
+var blockTrees = sync.Map{}
+
+type btSlice struct {
+	data    map[string]*BlockTree
+	changed time.Time
+	m       *sync.Mutex
+}
 
 type BlockTree struct {
 	ID       string // 块 ID
@@ -49,72 +55,61 @@ type BlockTree struct {
 }
 
 func GetRootUpdated() (ret map[string]string) {
-	blockTreesLock.Lock()
-	defer blockTreesLock.Unlock()
-
 	ret = map[string]string{}
-	for _, b := range blockTrees {
-		if b.RootID == b.ID {
-			ret[b.RootID] = b.Updated
+	blockTrees.Range(func(key, value interface{}) bool {
+		slice := value.(*btSlice)
+		slice.m.Lock()
+		for _, b := range slice.data {
+			if b.RootID == b.ID {
+				ret[b.RootID] = b.Updated
+			}
 		}
-	}
+		slice.m.Unlock()
+		return true
+	})
 	return
 }
 
-func GetBlockTreeByPath(path string) *BlockTree {
-	blockTreesLock.Lock()
-	defer blockTreesLock.Unlock()
-
-	for _, b := range blockTrees {
-		if b.Path == path {
-			return b
+func GetBlockTreeByPath(path string) (ret *BlockTree) {
+	blockTrees.Range(func(key, value interface{}) bool {
+		slice := value.(*btSlice)
+		slice.m.Lock()
+		for _, b := range slice.data {
+			if b.Path == path {
+				ret = b
+				break
+			}
 		}
-	}
-	return nil
+		slice.m.Unlock()
+		return nil == ret
+	})
+	return
 }
 
 func CountTrees() (ret int) {
-	blockTreesLock.Lock()
-	defer blockTreesLock.Unlock()
-
 	roots := map[string]bool{}
-	for _, b := range blockTrees {
-		roots[b.RootID] = true
-	}
+	blockTrees.Range(func(key, value interface{}) bool {
+		slice := value.(*btSlice)
+		slice.m.Lock()
+		for _, b := range slice.data {
+			roots[b.RootID] = true
+		}
+		slice.m.Unlock()
+		return true
+	})
 	ret = len(roots)
 	return
 }
 
 func CountBlocks() (ret int) {
-	blockTreesLock.Lock()
-	defer blockTreesLock.Unlock()
-	return len(blockTrees)
-}
-
-func CeilTreeCount(count int) int {
-	if 100 > count {
-		return 100
-	}
-
-	for i := 1; i < 40; i++ {
-		if count < i*500 {
-			return i * 500
-		}
-	}
-	return 500*40 + 1
-}
-
-func CeilBlockCount(count int) int {
-	if 5000 > count {
-		return 5000
-	}
-
-	for i := 1; i < 100; i++ {
-		if count < i*10000 {
-			return i * 10000
-		}
-	}
-	return 10000*100 + 1
+	blockTrees.Range(func(key, value interface{}) bool {
+		slice := value.(*btSlice)
+		slice.m.Lock()
+		ret += len(slice.data)
+		slice.m.Unlock()
+		return true
+	})
+	return
 }
 
 func GetRedundantPaths(boxID string, paths []string) (ret []string) {
@@ -123,15 +118,18 @@ func GetRedundantPaths(boxID string, paths []string) (ret []string) {
 		pathsMap[path] = true
 	}
 
-	tmp := blockTrees
 	btPathsMap := map[string]bool{}
-	for _, blockTree := range tmp {
-		if blockTree.BoxID != boxID {
-			continue
+	blockTrees.Range(func(key, value interface{}) bool {
+		slice := value.(*btSlice)
+		slice.m.Lock()
+		for _, b := range slice.data {
+			if b.BoxID == boxID {
+				btPathsMap[b.Path] = true
+			}
 		}
-
-		btPathsMap[blockTree.Path] = true
-	}
+		slice.m.Unlock()
+		return true
+	})
 
 	for p, _ := range btPathsMap {
 		if !pathsMap[p] {
@@ -148,15 +146,18 @@ func GetNotExistPaths(boxID string, paths []string) (ret []string) {
 		pathsMap[path] = true
 	}
 
-	tmp := blockTrees
 	btPathsMap := map[string]bool{}
-	for _, blockTree := range tmp {
-		if blockTree.BoxID != boxID {
-			continue
+	blockTrees.Range(func(key, value interface{}) bool {
+		slice := value.(*btSlice)
+		slice.m.Lock()
+		for _, b := range slice.data {
+			if b.BoxID == boxID {
+				btPathsMap[b.Path] = true
+			}
 		}
-
-		btPathsMap[blockTree.Path] = true
-	}
+		slice.m.Unlock()
+		return true
+	})
 
 	for p, _ := range pathsMap {
 		if !btPathsMap[p] {
@@ -167,132 +168,203 @@ func GetNotExistPaths(boxID string, paths []string) (ret []string) {
 	return
 }
 
-func GetBlockTreeRootByPath(boxID, path string) *BlockTree {
-	blockTreesLock.Lock()
-	defer blockTreesLock.Unlock()
-
-	for _, blockTree := range blockTrees {
-		if blockTree.BoxID == boxID && blockTree.Path == path && blockTree.RootID == blockTree.ID {
-			return blockTree
+func GetBlockTreeRootByPath(boxID, path string) (ret *BlockTree) {
+	blockTrees.Range(func(key, value interface{}) bool {
+		slice := value.(*btSlice)
+		slice.m.Lock()
+		for _, b := range slice.data {
+			if b.BoxID == boxID && b.Path == path && b.RootID == b.ID {
+				ret = b
+				break
+			}
 		}
-	}
-	return nil
+		slice.m.Unlock()
+		return nil == ret
+	})
+	return
 }
 
-func GetBlockTreeRootByHPath(boxID, hPath string) *BlockTree {
-	blockTreesLock.Lock()
-	defer blockTreesLock.Unlock()
-
-	for _, blockTree := range blockTrees {
-		if blockTree.BoxID == boxID && blockTree.HPath == hPath && blockTree.RootID == blockTree.ID {
-			return blockTree
+func GetBlockTreeRootByHPath(boxID, hPath string) (ret *BlockTree) {
+	blockTrees.Range(func(key, value interface{}) bool {
+		slice := value.(*btSlice)
+		slice.m.Lock()
+		for _, b := range slice.data {
+			if b.BoxID == boxID && b.HPath == hPath && b.RootID == b.ID {
+				ret = b
+				break
+			}
 		}
-	}
-	return nil
+		slice.m.Unlock()
+		return nil == ret
+	})
+	return
 }
 
-func GetBlockTree(id string) *BlockTree {
+func GetBlockTree(id string) (ret *BlockTree) {
 	if "" == id {
-		return nil
+		return
 	}
 
-	blockTreesLock.Lock()
-	defer blockTreesLock.Unlock()
-	return blockTrees[id]
+	hash := btHash(id)
+	val, ok := blockTrees.Load(hash)
+	if !ok {
+		return
+	}
+	slice := val.(*btSlice)
+	slice.m.Lock()
+	ret = slice.data[id]
+	slice.m.Unlock()
+	return
 }
 
 func SetBlockTreePath(tree *parse.Tree) {
-	blockTreesLock.Lock()
-	defer blockTreesLock.Unlock()
-
-	for _, b := range blockTrees {
-		if b.RootID == tree.ID {
-			b.BoxID, b.Path, b.HPath, b.Updated, b.Type = tree.Box, tree.Path, tree.HPath, tree.Root.IALAttr("updated"), TypeAbbr(ast.NodeDocument.String())
-		}
-	}
-	blockTreesChanged = time.Now()
+	hash := btHash(tree.ID)
+	val, ok := blockTrees.Load(hash)
+	if !ok {
+		val = &btSlice{data: map[string]*BlockTree{}, changed: time.Time{}, m: &sync.Mutex{}}
+		blockTrees.Store(hash, val)
+	}
+
+	slice := val.(*btSlice)
+	slice.m.Lock()
+	slice.data[tree.ID] = &BlockTree{
+		ID:      tree.ID,
+		RootID:  tree.Root.ID,
+		BoxID:   tree.Box,
+		Path:    tree.Path,
+		HPath:   tree.HPath,
+		Updated: tree.Root.IALAttr("updated"),
+		Type:    TypeAbbr(ast.NodeDocument.String()),
+	}
+	slice.m.Unlock()
+	slice.changed = time.Now()
 }
 
 func RemoveBlockTreesByRootID(rootID string) {
-	blockTreesLock.Lock()
-	defer blockTreesLock.Unlock()
-
 	var ids []string
-	for _, b := range blockTrees {
-		if b.RootID == rootID {
-			ids = append(ids, b.RootID)
+	blockTrees.Range(func(key, value interface{}) bool {
+		slice := value.(*btSlice)
+		slice.m.Lock()
+		for _, b := range slice.data {
+			if b.RootID == rootID {
+				ids = append(ids, b.RootID)
+			}
 		}
-	}
+		slice.m.Unlock()
+		return true
+	})
+
 	ids = gulu.Str.RemoveDuplicatedElem(ids)
 	for _, id := range ids {
-		delete(blockTrees, id)
+		val, ok := blockTrees.Load(btHash(id))
+		if !ok {
+			continue
+		}
+		slice := val.(*btSlice)
+		slice.m.Lock()
+		delete(slice.data, id)
+		slice.m.Unlock()
+		slice.changed = time.Now()
 	}
-	blockTreesChanged = time.Now()
 }
 
 func RemoveBlockTreesByPath(path string) {
-	blockTreesLock.Lock()
-	defer blockTreesLock.Unlock()
-
 	var ids []string
-	for _, b := range blockTrees {
-		if b.Path == path {
-			ids = append(ids, b.ID)
+	blockTrees.Range(func(key, value interface{}) bool {
+		slice := value.(*btSlice)
+		slice.m.Lock()
+		for _, b := range slice.data {
+			if b.Path == path {
+				ids = append(ids, b.RootID)
+			}
 		}
-	}
+		slice.m.Unlock()
+		return true
+	})
+
 	ids = gulu.Str.RemoveDuplicatedElem(ids)
 	for _, id := range ids {
-		delete(blockTrees, id)
+		val, ok := blockTrees.Load(btHash(id))
+		if !ok {
+			continue
+		}
+		slice := val.(*btSlice)
+		slice.m.Lock()
+		delete(slice.data, id)
+		slice.m.Unlock()
+		slice.changed = time.Now()
 	}
-	blockTreesChanged = time.Now()
 }
 
 func RemoveBlockTreesByPathPrefix(pathPrefix string) {
-	blockTreesLock.Lock()
-	defer blockTreesLock.Unlock()
-
 	var ids []string
-	for _, b := range blockTrees {
-		if strings.HasPrefix(b.Path, pathPrefix) {
-			ids = append(ids, b.ID)
+	blockTrees.Range(func(key, value interface{}) bool {
+		slice := value.(*btSlice)
+		slice.m.Lock()
+		for _, b := range slice.data {
+			if strings.HasPrefix(b.Path, pathPrefix) {
+				ids = append(ids, b.RootID)
+			}
 		}
-	}
+		slice.m.Unlock()
+		return true
+	})
+
 	ids = gulu.Str.RemoveDuplicatedElem(ids)
 	for _, id := range ids {
-		delete(blockTrees, id)
+		val, ok := blockTrees.Load(btHash(id))
+		if !ok {
+			continue
+		}
+		slice := val.(*btSlice)
+		slice.m.Lock()
+		delete(slice.data, id)
+		slice.m.Unlock()
+		slice.changed = time.Now()
 	}
-	blockTreesChanged = time.Now()
 }
 
 func RemoveBlockTreesByBoxID(boxID string) (ids []string) {
-	blockTreesLock.Lock()
-	defer blockTreesLock.Unlock()
-
-	for _, b := range blockTrees {
-		if b.BoxID == boxID {
-			ids = append(ids, b.ID)
+	blockTrees.Range(func(key, value interface{}) bool {
+		slice := value.(*btSlice)
+		slice.m.Lock()
+		for _, b := range slice.data {
+			if b.BoxID == boxID {
+				ids = append(ids, b.RootID)
+			}
 		}
-	}
+		slice.m.Unlock()
+		return true
+	})
+
 	ids = gulu.Str.RemoveDuplicatedElem(ids)
 	for _, id := range ids {
-		delete(blockTrees, id)
+		val, ok := blockTrees.Load(btHash(id))
+		if !ok {
+			continue
+		}
+		slice := val.(*btSlice)
+		slice.m.Lock()
+		delete(slice.data, id)
+		slice.m.Unlock()
+		slice.changed = time.Now()
 	}
-	blockTreesChanged = time.Now()
 	return
 }
 
 func RemoveBlockTree(id string) {
-	blockTreesLock.Lock()
-	defer blockTreesLock.Unlock()
-
-	delete(blockTrees, id)
-	blockTreesChanged = time.Now()
+	val, ok := blockTrees.Load(btHash(id))
+	if !ok {
+		return
+	}
+	slice := val.(*btSlice)
+	slice.m.Lock()
+	delete(slice.data, id)
+	slice.m.Unlock()
+	slice.changed = time.Now()
 }
 
 func IndexBlockTree(tree *parse.Tree) {
-	blockTreesLock.Lock()
-	defer blockTreesLock.Unlock()
-
 	ast.Walk(tree.Root, func(n *ast.Node, entering bool) ast.WalkStatus {
 		if !entering || !n.IsBlock() {
 			return ast.WalkContinue
@@ -304,10 +376,27 @@ func IndexBlockTree(tree *parse.Tree) {
 		if "" == n.ID {
 			return ast.WalkContinue
 		}
-		blockTrees[n.ID] = &BlockTree{ID: n.ID, ParentID: parentID, RootID: tree.ID, BoxID: tree.Box, Path: tree.Path, HPath: tree.HPath, Updated: tree.Root.IALAttr("updated"), Type: TypeAbbr(n.Type.String())}
+
+		hash := btHash(n.ID)
+		val, ok := blockTrees.Load(hash)
+		if !ok {
+			val = &btSlice{data: map[string]*BlockTree{}, changed: time.Time{}, m: &sync.Mutex{}}
+			blockTrees.Store(hash, val)
+		}
+		slice := val.(*btSlice)
+		slice.m.Lock()
+		if bt := slice.data[n.ID]; nil != bt {
+			if bt.Updated != n.IALAttr("updated") {
+				slice.data[n.ID] = &BlockTree{ID: n.ID, ParentID: parentID, RootID: tree.ID, BoxID: tree.Box, Path: tree.Path, HPath: tree.HPath, Updated: n.IALAttr("updated"), Type: TypeAbbr(n.Type.String())}
+				slice.changed = time.Now()
+			}
+		} else {
+			slice.data[n.ID] = &BlockTree{ID: n.ID, ParentID: parentID, RootID: tree.ID, BoxID: tree.Box, Path: tree.Path, HPath: tree.HPath, Updated: n.IALAttr("updated"), Type: TypeAbbr(n.Type.String())}
+			slice.changed = time.Now()
+		}
+		slice.m.Unlock()
 		return ast.WalkContinue
 	})
-	blockTreesChanged = time.Now()
 }
 
 func AutoFlushBlockTree() {
@@ -328,68 +417,129 @@ func InitBlockTree(force bool) {
 		return
 	}
 
-	var err error
-	fh, err := os.OpenFile(util.BlockTreePath, os.O_RDWR, 0644)
+	entries, err := os.ReadDir(util.BlockTreePath)
 	if nil != err {
-		logging.LogErrorf("open block tree file failed: %s", err)
+		logging.LogErrorf("read block tree dir failed: %s", err)
 		os.Exit(util.ExitCodeBlockTreeErr)
 		return
 	}
-	defer fh.Close()
 
-	data, err := io.ReadAll(fh)
-	if nil != err {
-		logging.LogErrorf("read block tree failed: %s", err)
-		os.Exit(util.ExitCodeBlockTreeErr)
-		return
-	}
-	blockTreesLock.Lock()
-	if err = msgpack.Unmarshal(data, &blockTrees); nil != err {
-		logging.LogErrorf("unmarshal block tree failed: %s", err)
-		if err = os.RemoveAll(util.BlockTreePath); nil != err {
-			logging.LogErrorf("removed corrupted block tree failed: %s", err)
+	size := uint64(0)
+	for _, entry := range entries {
+		if !strings.HasSuffix(entry.Name(), ".msgpack") {
+			continue
 		}
-		os.Exit(util.ExitCodeBlockTreeErr)
-		return
+
+		p := filepath.Join(util.BlockTreePath, entry.Name())
+		var fh *os.File
+		fh, err = os.OpenFile(p, os.O_RDWR, 0644)
+		if nil != err {
+			logging.LogErrorf("open block tree file failed: %s", err)
+			os.Exit(util.ExitCodeBlockTreeErr)
+			return
+		}
+
+		var data []byte
+		data, err = io.ReadAll(fh)
+		fh.Close()
+		if nil != err {
+			logging.LogErrorf("read block tree failed: %s", err)
+			os.Exit(util.ExitCodeBlockTreeErr)
+			return
+		}
+
+		sliceData := map[string]*BlockTree{}
+		if err = msgpack.Unmarshal(data, &sliceData); nil != err {
+			logging.LogErrorf("unmarshal block tree failed: %s", err)
+			if err = os.RemoveAll(util.BlockTreePath); nil != err {
+				logging.LogErrorf("removed corrupted block tree failed: %s", err)
+			}
+			os.Exit(util.ExitCodeBlockTreeErr)
+			return
+		}
+
+		name := entry.Name()[0:strings.Index(entry.Name(), ".")]
+		blockTrees.Store(name, &btSlice{data: sliceData, changed: time.Time{}, m: &sync.Mutex{}})
+		size += uint64(len(data))
 	}
-	blockTreesLock.Unlock()
+
 	runtime.GC()
 
 	if elapsed := time.Since(start).Seconds(); 2 < elapsed {
-		logging.LogWarnf("read block tree [%s] to [%s], elapsed [%.2fs]", humanize.Bytes(uint64(len(data))), util.BlockTreePath, elapsed)
+		logging.LogWarnf("read block tree [%s] to [%s], elapsed [%.2fs]", humanize.Bytes((size)), util.BlockTreePath, elapsed)
 	}
 	return
 }
 
 func SaveBlockTree(force bool) {
-	if !force && blockTreesChanged.IsZero() {
+	if force {
+
 		return
 	}
 
 	start := time.Now()
-	if blockTreesChanged.After(start.Add(-7 * time.Second)) {
-		return
+	os.MkdirAll(util.BlockTreePath, 0755)
+
+	size := uint64(0)
+	blockTrees.Range(func(key, value interface{}) bool {
+		slice := value.(*btSlice)
+		if !force && (slice.changed.IsZero() || slice.changed.After(start.Add(-7*time.Second))) {
+			return true
+		}
+
+		slice.m.Lock()
+		data, err := msgpack.Marshal(slice.data)
+		if nil != err {
+			logging.LogErrorf("marshal block tree failed: %s", err)
+			os.Exit(util.ExitCodeBlockTreeErr)
+			return false
+		}
+		slice.m.Unlock()
+
+		p := filepath.Join(util.BlockTreePath, key.(string)) + ".msgpack"
+		if err = gulu.File.WriteFileSafer(p, data, 0644); nil != err {
+			logging.LogErrorf("write block tree failed: %s", err)
+			os.Exit(util.ExitCodeBlockTreeErr)
+			return false
+		}
+		slice.changed = time.Time{}
+		size += uint64(len(data))
+		return true
+	})
+
+	runtime.GC()
+
+	if elapsed := time.Since(start).Seconds(); 2 < elapsed {
+		logging.LogWarnf("save block tree [size=%s] to [%s], elapsed [%.2fs]", humanize.Bytes(size), util.BlockTreePath, elapsed)
 	}
+}
 
-	blockTreesLock.Lock()
-	data, err := msgpack.Marshal(blockTrees)
-	if nil != err {
-		logging.LogErrorf("marshal block tree failed: %s", err)
-		os.Exit(util.ExitCodeBlockTreeErr)
-		return
+func CeilTreeCount(count int) int {
+	if 100 > count {
+		return 100
 	}
-	blockTreesLock.Unlock()
 
-	if err = gulu.File.WriteFileSafer(util.BlockTreePath, data, 0644); nil != err {
-		logging.LogErrorf("write block tree failed: %s", err)
-		os.Exit(util.ExitCodeBlockTreeErr)
-		return
+	for i := 1; i < 40; i++ {
+		if count < i*500 {
+			return i * 500
+		}
 	}
-	runtime.GC()
+	return 500*40 + 1
+}
 
-	if elapsed := time.Since(start).Seconds(); 2 < elapsed {
-		logging.LogWarnf("save block tree [size=%s] to [%s], elapsed [%.2fs]", humanize.Bytes(uint64(len(data))), util.BlockTreePath, elapsed)
+func CeilBlockCount(count int) int {
+	if 5000 > count {
+		return 5000
 	}
 
-	blockTreesChanged = time.Time{}
+	for i := 1; i < 100; i++ {
+		if count < i*10000 {
+			return i * 10000
+		}
+	}
+	return 10000*100 + 1
+}
+
+func btHash(id string) string {
+	return util2.Hash([]byte(id))[0:2]
 }

+ 1 - 1
kernel/util/working.go

@@ -253,7 +253,7 @@ func initWorkspaceDir(workspaceArg string) {
 	os.Setenv("TMP", osTmpDir)
 	DBPath = filepath.Join(TempDir, DBName)
 	HistoryDBPath = filepath.Join(TempDir, "history.db")
-	BlockTreePath = filepath.Join(TempDir, "blocktree.msgpack")
+	BlockTreePath = filepath.Join(TempDir, "blocktree")
 	SnippetsPath = filepath.Join(DataDir, "snippets")
 }
 

+ 1 - 1
kernel/util/working_mobile.go

@@ -155,7 +155,7 @@ func initWorkspaceDirMobile(workspaceBaseDir string) {
 	os.Setenv("TMP", osTmpDir)
 	DBPath = filepath.Join(TempDir, DBName)
 	HistoryDBPath = filepath.Join(TempDir, "history.db")
-	BlockTreePath = filepath.Join(TempDir, "blocktree.msgpack")
+	BlockTreePath = filepath.Join(TempDir, "blocktree")
 	SnippetsPath = filepath.Join(DataDir, "snippets")
 
 	AppearancePath = filepath.Join(ConfDir, "appearance")