siyuan/kernel/treenode/blocktree.go

// SiYuan - Refactor your thinking
// Copyright (c) 2020-present, b3log.org
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.

package treenode

import (
	"os"
	"path/filepath"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/88250/go-humanize"
	"github.com/88250/gulu"
	"github.com/88250/lute/ast"
	"github.com/88250/lute/parse"
	"github.com/panjf2000/ants/v2"
	util2 "github.com/siyuan-note/dejavu/util"
	"github.com/siyuan-note/logging"
	"github.com/siyuan-note/siyuan/kernel/task"
	"github.com/siyuan-note/siyuan/kernel/util"
	"github.com/vmihailenco/msgpack/v5"
)

var blockTrees = &sync.Map{}

type btSlice struct {
	data    map[string]*BlockTree
	changed time.Time
	m       *sync.Mutex
}

type BlockTree struct {
	ID       string // 块 ID
	RootID   string // 根 ID
	ParentID string // 父 ID
	BoxID    string // 笔记本 ID
	Path     string // 文档数据路径
	HPath    string // 文档可读路径
	Updated  string // 更新时间
	Type     string // 类型
}

func GetBlockTreesByType(typ string) (ret []*BlockTree) {
	blockTrees.Range(func(key, value interface{}) bool {
		slice := value.(*btSlice)
		slice.m.Lock()
		for _, b := range slice.data {
			if b.Type == typ {
				ret = append(ret, b)
			}
		}
		slice.m.Unlock()
		return true
	})
	return
}

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
			}
		}
		slice.m.Unlock()
		return nil == ret
	})
	return
}

func CountTrees() (ret int) {
	roots := map[string]bool{}
	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) {
	blockTrees.Range(func(key, value interface{}) bool {
		slice := value.(*btSlice)
		slice.m.Lock()
		ret += len(slice.data)
		slice.m.Unlock()
		return true
	})
	return
}

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
			}
		}
		slice.m.Unlock()
		return nil == ret
	})
	return
}

func GetBlockTreeRootByHPath(boxID, hPath string) (ret *BlockTree) {
	hPath = gulu.Str.RemoveInvisible(hPath)
	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
			}
		}
		slice.m.Unlock()
		return nil == ret
	})
	return
}

func GetBlockTreeRootsByHPath(boxID, hPath string) (ret []*BlockTree) {
	hPath = gulu.Str.RemoveInvisible(hPath)
	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 = append(ret, b)
			}
		}
		slice.m.Unlock()
		return true
	})
	return
}

func GetBlockTreeRootByHPathPreferredParentID(boxID, hPath, preferredParentID string) (ret *BlockTree) {
	hPath = gulu.Str.RemoveInvisible(hPath)
	var roots []*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 {
				if "" == preferredParentID {
					ret = b
					break
				}

				roots = append(roots, b)
			}
		}
		slice.m.Unlock()
		return nil == ret
	})
	if 1 > len(roots) {
		return
	}

	for _, root := range roots {
		if root.ID == preferredParentID {
			ret = root
			return
		}
	}
	ret = roots[0]
	return
}

func ExistBlockTree(id string) bool {
	hash := btHash(id)
	val, ok := blockTrees.Load(hash)
	if !ok {
		return false
	}
	slice := val.(*btSlice)
	slice.m.Lock()
	_, ok = slice.data[id]
	slice.m.Unlock()
	return ok
}

func GetBlockTree(id string) (ret *BlockTree) {
	if "" == id {
		return
	}

	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) {
	RemoveBlockTreesByRootID(tree.ID)
	IndexBlockTree(tree)
}

func RemoveBlockTreesByRootID(rootID string) {
	var ids []string
	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.ID)
			}
		}
		slice.m.Unlock()
		return true
	})

	ids = gulu.Str.RemoveDuplicatedElem(ids)
	for _, id := range ids {
		val, ok := blockTrees.Load(btHash(id))
		if !ok {
			continue
		}
		slice := val.(*btSlice)
		slice.m.Lock()
		delete(slice.data, id)
		slice.changed = time.Now()
		slice.m.Unlock()
	}
}

func GetBlockTreesByPathPrefix(pathPrefix string) (ret []*BlockTree) {
	blockTrees.Range(func(key, value interface{}) bool {
		slice := value.(*btSlice)
		slice.m.Lock()
		for _, b := range slice.data {
			if strings.HasPrefix(b.Path, pathPrefix) {
				ret = append(ret, b)
			}
		}
		slice.m.Unlock()
		return true
	})
	return
}

func GetBlockTreesByRootID(rootID string) (ret []*BlockTree) {
	blockTrees.Range(func(key, value interface{}) bool {
		slice := value.(*btSlice)
		slice.m.Lock()
		for _, b := range slice.data {
			if b.RootID == rootID {
				ret = append(ret, b)
			}
		}
		slice.m.Unlock()
		return true
	})
	return
}

func RemoveBlockTreesByPathPrefix(pathPrefix string) {
	var ids []string
	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.ID)
			}
		}
		slice.m.Unlock()
		return true
	})

	ids = gulu.Str.RemoveDuplicatedElem(ids)
	for _, id := range ids {
		val, ok := blockTrees.Load(btHash(id))
		if !ok {
			continue
		}
		slice := val.(*btSlice)
		slice.m.Lock()
		delete(slice.data, id)
		slice.changed = time.Now()
		slice.m.Unlock()
	}
}

func GetBlockTreesByBoxID(boxID 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 {
				ret = append(ret, b)
			}
		}
		slice.m.Unlock()
		return true
	})
	return
}

func RemoveBlockTreesByBoxID(boxID string) (ids []string) {
	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.ID)
			}
		}
		slice.m.Unlock()
		return true
	})

	ids = gulu.Str.RemoveDuplicatedElem(ids)
	for _, id := range ids {
		val, ok := blockTrees.Load(btHash(id))
		if !ok {
			continue
		}
		slice := val.(*btSlice)
		slice.m.Lock()
		delete(slice.data, id)
		slice.changed = time.Now()
		slice.m.Unlock()
	}
	return
}

func RemoveBlockTree(id string) {
	val, ok := blockTrees.Load(btHash(id))
	if !ok {
		return
	}
	slice := val.(*btSlice)
	slice.m.Lock()
	delete(slice.data, id)
	slice.changed = time.Now()
	slice.m.Unlock()
}

func IndexBlockTree(tree *parse.Tree) {
	var changedNodes []*ast.Node
	ast.Walk(tree.Root, func(n *ast.Node, entering bool) ast.WalkStatus {
		if !entering || !n.IsBlock() {
			return ast.WalkContinue
		}
		if "" == n.ID {
			return ast.WalkContinue
		}

		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()
		bt := slice.data[n.ID]
		slice.m.Unlock()

		if nil != bt {
			if bt.Updated != n.IALAttr("updated") || bt.Type != TypeAbbr(n.Type.String()) || bt.Path != tree.Path || bt.BoxID != tree.Box || bt.HPath != tree.HPath {
				children := ChildBlockNodes(n) // 需要考虑子块，因为一些操作（比如移动块）后需要同时更新子块
				changedNodes = append(changedNodes, children...)
			}
		} else {
			children := ChildBlockNodes(n)
			changedNodes = append(changedNodes, children...)
		}
		return ast.WalkContinue
	})

	for _, n := range changedNodes {
		updateBtSlice(n, tree)
	}
}

func updateBtSlice(n *ast.Node, tree *parse.Tree) {
	var parentID string
	if nil != n.Parent {
		parentID = n.Parent.ID
	}

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

var blockTreeLock = sync.Mutex{}

func InitBlockTree(force bool) {
	blockTreeLock.Lock()
	defer blockTreeLock.Unlock()

	start := time.Now()
	if force {
		err := os.RemoveAll(util.BlockTreePath)
		if nil != err {
			logging.LogErrorf("remove block tree file failed: %s", err)
		}
		blockTrees = &sync.Map{}
		return
	}

	entries, err := os.ReadDir(util.BlockTreePath)
	if nil != err {
		logging.LogErrorf("read block tree dir failed: %s", err)
		os.Exit(logging.ExitCodeFileSysErr)
		return
	}

	loadErr := atomic.Bool{}
	size := atomic.Int64{}
	waitGroup := &sync.WaitGroup{}
	p, _ := ants.NewPoolWithFunc(4, func(arg interface{}) {
		defer waitGroup.Done()

		entry := arg.(os.DirEntry)
		p := filepath.Join(util.BlockTreePath, entry.Name())

		f, err := os.OpenFile(p, os.O_RDONLY, 0644)
		if nil != err {
			logging.LogErrorf("open block tree failed: %s", err)
			loadErr.Store(true)
			return
		}
		defer f.Close()

		info, err := f.Stat()
		if nil != err {
			logging.LogErrorf("stat block tree failed: %s", err)
			loadErr.Store(true)
			return
		}
		size.Add(info.Size())

		sliceData := map[string]*BlockTree{}
		if err = msgpack.NewDecoder(f).Decode(&sliceData); nil != err {
			logging.LogErrorf("unmarshal block tree failed: %s", err)
			loadErr.Store(true)
			return
		}

		name := entry.Name()[0:strings.Index(entry.Name(), ".")]
		blockTrees.Store(name, &btSlice{data: sliceData, changed: time.Time{}, m: &sync.Mutex{}})
	})
	for _, entry := range entries {
		if !strings.HasSuffix(entry.Name(), ".msgpack") {
			continue
		}

		waitGroup.Add(1)
		p.Invoke(entry)
	}

	waitGroup.Wait()
	p.Release()

	if loadErr.Load() {
		logging.LogInfof("cause block tree load error, remove block tree file")
		if removeErr := os.RemoveAll(util.BlockTreePath); nil != removeErr {
			logging.LogErrorf("remove block tree file failed: %s", removeErr)
			os.Exit(logging.ExitCodeFileSysErr)
			return
		}
		blockTrees = &sync.Map{}
		return
	}

	elapsed := time.Since(start).Seconds()
	logging.LogInfof("read block tree [%s] to [%s], elapsed [%.2fs]", humanize.BytesCustomCeil(uint64(size.Load()), 2), util.BlockTreePath, elapsed)
	return
}

func SaveBlockTreeJob() {
	SaveBlockTree(false)
}

func SaveBlockTree(force bool) {
	blockTreeLock.Lock()
	defer blockTreeLock.Unlock()

	if task.ContainIndexTask() {
		//logging.LogInfof("skip saving block tree because indexing")
		return
	}
	//logging.LogInfof("saving block tree")

	start := time.Now()
	if err := os.MkdirAll(util.BlockTreePath, 0755); nil != err {
		logging.LogErrorf("create block tree dir [%s] failed: %s", util.BlockTreePath, err)
		os.Exit(logging.ExitCodeFileSysErr)
		return
	}

	size := uint64(0)
	var count int
	blockTrees.Range(func(key, value interface{}) bool {
		slice := value.(*btSlice)
		slice.m.Lock()
		if !force && slice.changed.IsZero() {
			slice.m.Unlock()
			return true
		}

		data, err := msgpack.Marshal(slice.data)
		if nil != err {
			logging.LogErrorf("marshal block tree failed: %s", err)
			os.Exit(logging.ExitCodeFileSysErr)
			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(logging.ExitCodeFileSysErr)
			return false
		}

		slice.m.Lock()
		slice.changed = time.Time{}
		slice.m.Unlock()
		size += uint64(len(data))
		count++
		return true
	})
	if 0 < count {
		//logging.LogInfof("wrote block trees [%d]", count)
	}

	elapsed := time.Since(start).Seconds()
	if 2 < elapsed {
		logging.LogWarnf("save block tree [size=%s] to [%s], elapsed [%.2fs]", humanize.BytesCustomCeil(size, 2), util.BlockTreePath, elapsed)
	}
}

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
}

func btHash(id string) string {
	return util2.Hash([]byte(id))[0:2]
}