1
0
mirror of https://github.com/Mrs4s/go-cqhttp.git synced 2025-05-04 19:17:37 +08:00

feat: init disk-based btree impl

This commit is contained in:
wdvxdr 2021-09-27 22:44:45 +08:00
parent c4d34fa14f
commit f767213681
No known key found for this signature in database
GPG Key ID: 703F8C071DE7A1B6
14 changed files with 692 additions and 12 deletions

View File

@ -92,4 +92,4 @@ issues:
fix: true
exclude-use-default: false
exclude:
- "Error return value of .((os.)?std(out|err)..*|.*Close|.*Flush|os.Remove(All)?|.*print(f|ln)?|os.(Un)?Setenv). is not check"
- "Error return value of .((os.)?std(out|err)..*|.*Close|.*Seek|.*Flush|os.Remove(All)?|.*print(f|ln)?|os.(Un)?Setenv). is not check"

View File

@ -5,7 +5,6 @@ import (
"encoding/hex"
"encoding/json"
"fmt"
"github.com/Mrs4s/go-cqhttp/db"
"math"
"os"
"path"
@ -15,6 +14,8 @@ import (
"strings"
"time"
"github.com/Mrs4s/go-cqhttp/db"
"github.com/Mrs4s/MiraiGo/binary"
"github.com/Mrs4s/MiraiGo/client"
"github.com/Mrs4s/MiraiGo/message"

View File

@ -5,7 +5,6 @@ import (
"encoding/hex"
"encoding/json"
"fmt"
"github.com/Mrs4s/go-cqhttp/db"
"io"
"os"
"path"
@ -13,6 +12,8 @@ import (
"sync"
"time"
"github.com/Mrs4s/go-cqhttp/db"
"github.com/Mrs4s/MiraiGo/binary"
"github.com/Mrs4s/MiraiGo/client"
"github.com/Mrs4s/MiraiGo/message"

View File

@ -7,7 +7,6 @@ import (
xml2 "encoding/xml"
"errors"
"fmt"
"github.com/Mrs4s/go-cqhttp/db"
"io"
"math/rand"
"net/url"
@ -18,6 +17,8 @@ import (
"strings"
"time"
"github.com/Mrs4s/go-cqhttp/db"
"github.com/Mrs4s/MiraiGo/binary"
"github.com/Mrs4s/MiraiGo/message"
"github.com/Mrs4s/MiraiGo/utils"

View File

@ -2,13 +2,14 @@ package coolq
import (
"encoding/hex"
"github.com/Mrs4s/go-cqhttp/db"
"os"
"path"
"strconv"
"strings"
"time"
"github.com/Mrs4s/go-cqhttp/db"
"github.com/Mrs4s/go-cqhttp/global"
"github.com/Mrs4s/go-cqhttp/internal/base"

View File

@ -2,8 +2,9 @@ package db
import (
"fmt"
"github.com/Mrs4s/go-cqhttp/global"
"hash/crc32"
"github.com/Mrs4s/go-cqhttp/global"
)
type (

View File

@ -3,12 +3,14 @@ package db
import (
"bytes"
"encoding/gob"
"path"
"github.com/Mrs4s/MiraiGo/binary"
"github.com/Mrs4s/go-cqhttp/global"
"github.com/pkg/errors"
"github.com/syndtr/goleveldb/leveldb"
"github.com/syndtr/goleveldb/leveldb/opt"
"path"
"github.com/Mrs4s/go-cqhttp/global"
)
type LevelDBImpl struct {

View File

@ -2,6 +2,7 @@ package db
import (
"context"
"github.com/pkg/errors"
"go.mongodb.org/mongo-driver/bson"
"go.mongodb.org/mongo-driver/mongo"

1
go.mod
View File

@ -51,7 +51,6 @@ require (
github.com/xdg-go/scram v1.0.2 // indirect
github.com/xdg-go/stringprep v1.0.2 // indirect
github.com/youmark/pkcs8 v0.0.0-20181117223130-1be2e3e5546d // indirect
golang.org/x/image v0.0.0-20210628002857-a66eb6448b8d // indirect
golang.org/x/net v0.0.0-20210505024714-0287a6fb4125 // indirect
golang.org/x/sync v0.0.0-20210220032951-036812b2e83c // indirect
golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1 // indirect

2
go.sum
View File

@ -178,8 +178,6 @@ golang.org/x/crypto v0.0.0-20200302210943-78000ba7a073/go.mod h1:LzIPMQfyMNhhGPh
golang.org/x/crypto v0.0.0-20210817164053-32db794688a5 h1:HWj/xjIHfjYU5nVXpTM0s39J9CbLn7Cc5a7IC5rwsMQ=
golang.org/x/crypto v0.0.0-20210817164053-32db794688a5/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc=
golang.org/x/exp v0.0.0-20190121172915-509febef88a4/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA=
golang.org/x/image v0.0.0-20210628002857-a66eb6448b8d h1:RNPAfi2nHY7C2srAV8A49jpsYr0ADedCk1wq6fTMTvs=
golang.org/x/image v0.0.0-20210628002857-a66eb6448b8d/go.mod h1:023OzeP/+EPmXeapQh35lcL3II3LrY8Ic+EFFKVhULM=
golang.org/x/lint v0.0.0-20181026193005-c67002cb31c3/go.mod h1:UVdnD1Gm6xHRNCYTkRU2/jEulfH38KcIWyp/GAMgvoE=
golang.org/x/lint v0.0.0-20190227174305-5b3e6a55c961/go.mod h1:wehouNa3lNwaWXcvxsM5YxQ5yQlVC4a0KAMCusXpPoU=
golang.org/x/lint v0.0.0-20190313153728-d0100b6bd8b3/go.mod h1:6SW0HCj/g11FgYtHlgUYUwCkIfeOF89ocIRzGO/8vkc=

505
internal/btree/btree.go Normal file
View File

@ -0,0 +1,505 @@
// Package btree provide a disk-based btree
package btree
import (
"encoding/binary"
"io"
"math/rand"
"os"
"unsafe"
"github.com/pkg/errors"
)
const (
sha1Size = 20 // md5 sha1
tableSize = (4096 - 1) / int(unsafe.Sizeof(item{}))
cacheSlots = 23 // prime
superSize = unsafe.Sizeof(super{})
tableStructSize = int(unsafe.Sizeof(table{}))
)
type item struct {
sha1 [sha1Size]byte
offset int64
child int64
}
type table struct {
items [tableSize]item
size int
}
type cache struct {
table *table
offset int64
}
type super struct {
top int64
freeTop int64
alloc int64
}
// Btree ...
type Btree struct {
fd *os.File
top int64
freeTop int64
alloc int64
cache [23]cache
inAllocator bool
deleteLarger bool
}
func (bt *Btree) get(offset int64) *table {
assert(offset != 0)
// take from cache
slot := &bt.cache[offset%cacheSlots]
if slot.offset == offset {
return slot.table
}
table := new(table)
bt.fd.Seek(offset, io.SeekStart)
err := binary.Read(bt.fd, binary.LittleEndian, table) // todo(wdvxdr): efficient reading
if err != nil {
panic(errors.Wrap(err, "btree I/O error"))
}
return table
}
func (bt *Btree) put(t *table, offset int64) {
assert(offset != 0)
/* overwrite cache */
slot := &bt.cache[offset%cacheSlots]
slot.table = t
slot.offset = offset
}
func (bt *Btree) flush(t *table, offset int64) {
assert(offset != 0)
bt.fd.Seek(offset, io.SeekStart)
err := binary.Write(bt.fd, binary.LittleEndian, t)
if err != nil {
panic(errors.Wrap(err, "btree I/O error"))
}
bt.put(t, offset)
}
func (bt *Btree) flushSuper() {
bt.fd.Seek(0, io.SeekStart)
super := super{
top: bt.top,
freeTop: bt.freeTop,
alloc: bt.alloc,
}
err := binary.Write(bt.fd, binary.LittleEndian, super)
if err != nil {
panic(errors.Wrap(err, "btree I/O error"))
}
}
// Open opens an existed btree file
func Open(name string) (*Btree, error) {
btree := new(Btree)
fd, err := os.OpenFile(name, os.O_RDWR, 0o644)
if err != nil {
return nil, errors.Wrap(err, "btree open file failed")
}
btree.fd = fd
super := super{}
err = binary.Read(fd, binary.LittleEndian, &super)
btree.top = super.top
btree.freeTop = super.freeTop
btree.alloc = super.alloc
return btree, errors.Wrap(err, "btree read meta info failed")
}
// Create creates a database
func Create(name string) (*Btree, error) {
btree := new(Btree)
fd, err := os.OpenFile(name, os.O_RDWR|os.O_TRUNC|os.O_CREATE, 0o644)
if err != nil {
return nil, errors.Wrap(err, "btree open file failed")
}
btree.fd = fd
btree.alloc = int64(superSize)
btree.flushSuper()
return btree, nil
}
// Close closes the database
func (bt *Btree) Close() error {
err := bt.Close()
for i := 0; i < cacheSlots; i++ {
bt.cache[i] = cache{}
}
return errors.Wrap(err, "btree close failed")
}
func collapse(bt *Btree, offset int64) int64 {
table := bt.get(offset)
if table.size != 0 {
/* unable to collapse */
bt.put(table, offset)
return offset
}
ret := table.items[0].child
bt.put(table, offset)
/*
* WARNING: this is dangerous as the chunk is added to allocation tree
* before the references to it are removed!
*/
bt.freeChunk(offset, int(unsafe.Sizeof(table)))
return ret
}
// split a table. The pivot item is stored to 'sha1' and 'offset'.
// Returns offset to the new table.
func (bt *Btree) split(t *table, hash *byte, offset *int64) int64 {
copysha1(hash, &t.items[tableSize/2].sha1[0])
*offset = t.items[tableSize/2].offset
ntable := new(table)
ntable.size = t.size - tableSize/2 - 1
t.size = tableSize / 2
copy(ntable.items[:ntable.size+1], t.items[tableSize/2+1:])
noff := bt.allocChunk(tableStructSize)
bt.flush(ntable, noff)
// make sure data is written before a reference is added to it
_ = bt.fd.Sync()
return noff
}
// takeSmallest find and remove the smallest item from the given table. The key of the item
// is stored to 'sha1'. Returns offset to the item
func (bt *Btree) takeSmallest(toff int64, sha1 *byte) int64 {
table := bt.get(toff)
assert(table.size > 0)
var off int64
child := table.items[0].child
if child == 0 {
off = bt.remove(table, 0, sha1)
} else {
/* recursion */
off = bt.takeSmallest(child, sha1)
table.items[0].child = collapse(bt, child)
}
bt.flush(table, toff)
// make sure data is written before a reference is added to it
_ = bt.fd.Sync()
return off
}
// takeLargest find and remove the largest item from the given table. The key of the item
// is stored to 'sha1'. Returns offset to the item
func (bt *Btree) takeLargest(toff int64, sha1 *byte) int64 {
table := bt.get(toff)
assert(table.size > 0)
var off int64
child := table.items[table.size].child
if child == 0 {
off = bt.remove(table, table.size-1, sha1)
} else {
/* recursion */
off = bt.takeLargest(child, sha1)
table.items[table.size].child = collapse(bt, child)
}
bt.flush(table, toff)
// make sure data is written before a reference is added to it
_ = bt.fd.Sync()
return off
}
// remove an item in position 'i' from the given table. The key of the
// removed item is stored to 'sha1'. Returns offset to the item.
func (bt *Btree) remove(t *table, i int, sha1 *byte) int64 {
assert(i < t.size)
if sha1 != nil {
copysha1(sha1, &t.items[i].sha1[0])
}
offset := t.items[i].offset
lc := t.items[i].child
rc := t.items[i+1].child
if lc != 0 && rc != 0 {
/* replace the removed item by taking an item from one of the
child tables */
var noff int64
if rand.Int()&1 != 0 {
noff = bt.takeLargest(lc, &t.items[i].sha1[0])
t.items[i].child = collapse(bt, lc)
} else {
noff = bt.takeSmallest(rc, &t.items[i].sha1[0])
t.items[i+1].child = collapse(bt, rc)
}
t.items[i].child = noff
} else {
// memmove(&table->items[i], &table->items[i + 1],
// (table->size - i) * sizeof(struct btree_item));
// table->size--;
for j := i; j < t.size-i; j++ { // fuck you, go!
t.items[j] = t.items[j+1]
}
t.size--
if lc != 0 {
t.items[i].child = lc
} else {
t.items[i].child = rc
}
}
return offset
}
func (bt *Btree) insert(toff int64, sha1 *byte, data []byte, len int) int64 {
table := bt.get(toff)
assert(table.size < tableSize-1)
left, right := 0, table.size
for left < right {
mid := (right-left)>>1 + left
switch cmp := cmp(sha1, &table.items[mid].sha1[0]); {
case cmp == 0:
// already in the table
ret := table.items[mid].offset
bt.put(table, toff)
return ret
case cmp < 0:
right = mid
default:
left = mid + 1
}
}
i := left
var off, rc, ret int64
lc := table.items[i].child
if lc != 0 {
/* recursion */
ret = bt.insert(lc, sha1, data, len)
/* check if we need to split */
child := bt.get(lc)
if child.size < tableSize-1 {
/* nothing to do */
bt.put(table, toff)
bt.put(child, lc)
return ret
}
/* overwrites SHA-1 */
rc = bt.split(child, sha1, &off)
/* flush just in case changes happened */
bt.flush(child, lc)
// make sure data is written before a reference is added to it
_ = bt.fd.Sync()
} else {
off = bt.insertData(data, len)
ret = off
}
table.size++
// todo:
// memmove(&table->items[i + 1], &table->items[i],
// (table->size - i) * sizeof(struct btree_item));
copysha1(&table.items[i].sha1[0], sha1)
table.items[i].offset = off
table.items[i].child = lc
table.items[i+1].child = rc
bt.flush(table, toff)
return ret
}
func (bt *Btree) insertData(data []byte, size int) int64 {
if data == nil {
return int64(size)
}
assert(len(data) == size)
offset := bt.allocChunk(4 + len(data))
bt.fd.Seek(offset, io.SeekStart)
err := binary.Write(bt.fd, binary.LittleEndian, int32(len(data)))
if err != nil {
panic(errors.Wrap(err, "btree I/O error"))
}
_, err = bt.fd.Write(data)
if err != nil {
panic(errors.Wrap(err, "btree I/O error"))
}
// make sure data is written before a reference is added to it
_ = bt.fd.Sync()
return offset
}
// delete remove an item with key 'sha1' from the given table. The offset to the
// removed item is returned.
// Please note that 'sha1' is overwritten when called inside the allocator.
func (bt *Btree) delete(offset int64, hash *byte) int64 {
if offset == 0 {
return 0
}
table := bt.get(offset)
left, right := 0, table.size
for left < right {
i := (right-left)>>1 + left
switch cmp := cmp(hash, &table.items[i].sha1[0]); {
case cmp == 0:
// found
ret := bt.remove(table, i, hash)
bt.flush(table, offset)
return ret
case cmp < 0:
right = i
default:
left = i + 1
}
}
// not found - recursion
i := left
child := table.items[i].child
ret := bt.delete(child, hash)
if ret != 0 {
table.items[i].child = collapse(bt, child)
}
if ret == 0 && bt.deleteLarger && i < table.size {
ret = bt.remove(table, i, hash)
}
if ret != 0 {
/* flush just in case changes happened */
bt.flush(table, offset)
} else {
bt.put(table, offset)
}
return ret
}
func (bt *Btree) insertTopLevel(toff *int64, sha1 *byte, data []byte, len int) int64 {
var off, ret, rc int64
if *toff != 0 {
ret = bt.insert(*toff, sha1, data, len)
/* check if we need to split */
table := bt.get(*toff)
if table.size < tableSize-1 {
/* nothing to do */
bt.put(table, *toff)
return ret
}
rc = bt.split(table, sha1, &off)
bt.flush(table, *toff)
} else {
off = bt.insertData(data, len)
ret = off
}
/* create new top level table */
t := new(table)
t.size = 1
copysha1(&t.items[0].sha1[0], sha1)
t.items[0].offset = off
t.items[0].child = *toff
t.items[1].child = rc
ntoff := bt.allocChunk(tableStructSize)
bt.flush(t, ntoff)
*toff = ntoff
// make sure data is written before a reference is added to it
_ = bt.fd.Sync()
return ret
}
func (bt *Btree) lookup(toff int64, sha1 *byte) int64 {
if toff == 0 {
return 0
}
table := bt.get(toff)
left, right := 0, table.size
for left < right {
mid := (right-left)>>1 + left
switch cmp := cmp(sha1, &table.items[mid].sha1[0]); {
case cmp == 0:
// found
ret := table.items[mid].offset
bt.put(table, toff)
return ret
case cmp < 0:
right = mid
default:
left = mid + 1
}
}
i := left
child := table.items[i].child
bt.put(table, toff)
return bt.lookup(child, sha1)
}
// Insert a new item with key 'sha1' with the contents in 'data' to the
// database file.
func (bt *Btree) Insert(csha1 *byte, data []byte) {
/* SHA-1 must be in writable memory */
var sha1 [sha1Size]byte
copysha1(&sha1[0], csha1)
bt.insertTopLevel(&bt.top, &sha1[0], data, len(data))
freeQueued(bt)
bt.flushSuper()
}
// Get look up item with the given key 'sha1' in the database file. Length of the
// item is stored in 'len'. Returns a pointer to the contents of the item.
// The returned pointer should be released with free() after use.
func (bt *Btree) Get(sha1 *byte) []byte {
off := bt.lookup(bt.top, sha1)
if off == 0 {
return nil
}
bt.fd.Seek(off, io.SeekStart)
var length int32
err := binary.Read(bt.fd, binary.LittleEndian, &length)
if err != nil {
return nil
}
data := make([]byte, length)
n, err := io.ReadFull(bt.fd, data)
if err != nil {
return nil
}
return data[:n]
}
// Delete remove item with the given key 'sha1' from the database file.
func (bt *Btree) Delete(sha1 *byte) error {
return errors.New("impl me")
}

128
internal/btree/chunk.go Normal file
View File

@ -0,0 +1,128 @@
package btree
import (
"math/rand"
"unsafe"
)
type chunk struct {
offset int64
len int
}
const freeQueueLen = 64
// todo(wdvxdr): move this to btree?
var (
fqueue [freeQueueLen]chunk
fqueueLen = 0
)
func freeQueued(bt *Btree) {
for i := 0; i < fqueueLen; i++ {
chunk := &fqueue[i]
bt.freeChunk(chunk.offset, chunk.len)
}
fqueueLen = 0
}
func (bt *Btree) allocChunk(len int) int64 {
assert(len > 0)
len = power2(len)
var offset int64
if bt.inAllocator {
const i32s = unsafe.Sizeof(int32(0))
/* create fake size SHA-1 */
var sha1 [sha1Size]byte
p := unsafe.Pointer(&sha1[0])
*(*int32)(p) = -1 // *(uint32_t *) sha1 = -1;
*(*uint32)(unsafe.Add(p, i32s)) = uint32(len) // ((__be32 *) sha1)[1] = to_be32(len);
/* find free chunk with the larger or the same len/SHA-1 */
bt.inAllocator = true
bt.deleteLarger = true
offset = bt.delete(bt.freeTop, &sha1[0])
bt.deleteLarger = false
if offset != 0 {
assert(*(*int32)(p) == -1) // assert(*(uint32_t *) sha1 == (uint32_t) -1)
flen := int(*(*uint32)(unsafe.Add(p, i32s))) // size_t free_len = from_be32(((__be32 *) sha1)[1])
assert(power2(flen) == flen)
assert(flen >= len)
/* delete buddy information */
resetsha1(&sha1[0])
*(*int64)(p) = offset
buddyLen := bt.delete(bt.freeTop, &sha1[0])
assert(buddyLen == int64(len))
bt.freeTop = collapse(bt, bt.freeTop)
bt.inAllocator = false
/* free extra space at the end of the chunk */
for flen > len {
flen >>= 1
bt.freeChunk(offset+int64(flen), flen)
}
} else {
bt.inAllocator = false
}
}
if offset == 0 {
/* not found, allocate from the end of the file */
offset = bt.alloc
/* TODO: this wastes memory.. */
if offset&int64(len-1) != 0 {
offset += int64(len) - (offset & (int64(len) - 1))
}
bt.alloc = offset + int64(len)
}
bt.flushSuper()
// make sure the allocation tree is up-to-date before using the chunk
_ = bt.fd.Sync()
return offset
}
/* Mark a chunk as unused in the database file */
func (bt *Btree) freeChunk(offset int64, len int) {
assert(len > 0)
assert(offset != 0)
len = power2(len)
assert(offset&int64(len-1) == 0)
if bt.inAllocator {
chunk := &fqueue[fqueueLen]
fqueueLen++
chunk.offset = offset
chunk.len = len
return
}
/* create fake offset SHA-1 for buddy allocation */
var sha1 [sha1Size]byte
p := unsafe.Pointer(&sha1[0])
bt.inAllocator = true
const i32s = unsafe.Sizeof(int32(0))
/* add buddy information */
resetsha1(&sha1[0])
*(*int32)(p) = -1 // *(uint32_t *) sha1 = -1;
*(*uint32)(unsafe.Add(p, i32s)) = uint32(len) // ((__be32 *) sha1)[1] = to_be32(len);
*(*uint32)(unsafe.Add(p, i32s*2)) = rand.Uint32() /* to make SHA-1 unique */
*(*uint32)(unsafe.Add(p, i32s*3)) = rand.Uint32()
// insert_toplevel(btree, &btree->free_top, sha1, NULL, offset);
bt.insertTopLevel(&bt.freeTop, &sha1[0], nil, int(offset))
bt.inAllocator = false
bt.flushSuper()
// make sure the allocation tree is up-to-date before removing
// references to the chunk
_ = bt.fd.Sync()
}

42
internal/btree/helper.go Normal file
View File

@ -0,0 +1,42 @@
package btree
import "unsafe"
func assert(cond bool) {
if !cond {
panic("assert failed!")
}
}
// power2 returns a value that is greater or equal to 'val' and is power-of-two.
func power2(val int) int {
i := 1
for i < val {
i <<= 1
}
return i
}
// helpers for sha1
func cmp(a, b *byte) uint64 {
pa, pb := unsafe.Pointer(a), unsafe.Pointer(b)
if *(*uint64)(pa) != *(*uint64)(pb) {
return *(*uint64)(pa) - *(*uint64)(pb)
}
pa, pb = unsafe.Add(pa, 8), unsafe.Add(pb, 8)
if *(*uint64)(pa) != *(*uint64)(pb) {
return *(*uint64)(pa) - *(*uint64)(pb)
}
return uint64(*(*uint32)(unsafe.Add(pa, 8)) - *(*uint32)(unsafe.Add(pa, 8)))
}
func copysha1(dst *byte, src *byte) {
pa, pb := unsafe.Pointer(dst), unsafe.Pointer(src)
*(*[sha1Size]byte)(pa) = *(*[sha1Size]byte)(pb)
}
func resetsha1(sha1 *byte) {
p := unsafe.Pointer(sha1)
*(*[sha1Size]byte)(p) = [sha1Size]byte{}
}

View File

@ -43,7 +43,7 @@ func check(r io.ReadSeeker, list []string) (bool, string) {
return true, ""
}
_, _ = r.Seek(0, io.SeekStart)
defer r.Seek(0, io.SeekStart) // nolint
defer r.Seek(0, io.SeekStart)
t, err := mimetype.DetectReader(r)
if err != nil {
logrus.Debugf("扫描 Mime 时出现问题: %v", err)