mirror of
https://github.com/ethereum/go-ethereum.git
synced 2026-08-20 10:52:25 +00:00
ethdb: introduce vectordb for persistent sequential data storage
The vectordb provides a database for storing binary blobs with a sequential relationship to one other. This is specifically targetting blockchain data constructs that include block header, block bodies, a block's transaction receipts, and block number to hash mappings. A considerable amount of this code comes from the "freezer" prototyping work done in #17814. However it does contain a few modifications: 1. It is not append-only and immutable. A VectorDB can shrink or grow so that it can be used for storing all relevant blockchain data in the presence of chain reorganizations. 2. The metadata contained within the index file has been moved to a struct so that it is easier to add metadata (if needed). 3. Adds a metadata file for the database instance (to store the version). This extracts out the "freezer table" from the "freezer" prototyping done in #17814, with a few modifications. The design decision to keep all data in a single file was kept, although it will probably make sense to break that up down the road for several reasons outside the scope of this description. TODO in this PR: 1. Add the ability to repair a VectorDB instance. 2. Add in logging.
This commit is contained in:
parent
729bf365b5
commit
45ad0d7b1a
2 changed files with 876 additions and 0 deletions
462
ethdb/vectordb/vectordb.go
Normal file
462
ethdb/vectordb/vectordb.go
Normal file
|
|
@ -0,0 +1,462 @@
|
||||||
|
// Copyright 2019 The go-ethereum Authors
|
||||||
|
// This file is part of the go-ethereum library.
|
||||||
|
//
|
||||||
|
// The go-ethereum library is free software: you can redistribute it and/or modify
|
||||||
|
// it under the terms of the GNU Lesser General Public License as published by
|
||||||
|
// the Free Software Foundation, either version 3 of the License, or
|
||||||
|
// (at your option) any later version.
|
||||||
|
//
|
||||||
|
// The go-ethereum library 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 Lesser General Public License for more details.
|
||||||
|
//
|
||||||
|
// You should have received a copy of the GNU Lesser General Public License
|
||||||
|
// along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.
|
||||||
|
|
||||||
|
// Package vectordb provides the vector database implementation.
|
||||||
|
package vectordb
|
||||||
|
|
||||||
|
import (
|
||||||
|
"encoding/binary"
|
||||||
|
"encoding/json"
|
||||||
|
"errors"
|
||||||
|
"fmt"
|
||||||
|
"io/ioutil"
|
||||||
|
"os"
|
||||||
|
"path/filepath"
|
||||||
|
"sync"
|
||||||
|
)
|
||||||
|
|
||||||
|
const (
|
||||||
|
// The current database version.
|
||||||
|
currentVersion = 1
|
||||||
|
|
||||||
|
// The file used to store the database metadata.
|
||||||
|
metadataFile = "METADATA"
|
||||||
|
// The file used to store the rawIndex of items stored in the database.
|
||||||
|
indexFile = "INDEX"
|
||||||
|
// The file used to store the rawData contained in the database.
|
||||||
|
dataFile = "DATA"
|
||||||
|
|
||||||
|
// Size of a serialized rawIndex entry. This should be updated if any changes
|
||||||
|
// are made to indexEntry.
|
||||||
|
indexEntryLen = 16 // 2 * sizeof(uint64)
|
||||||
|
|
||||||
|
// File permissions.
|
||||||
|
dbDirPerm = 0755
|
||||||
|
indexFilePerm = 0644
|
||||||
|
dataFilePerm = 0644
|
||||||
|
metadataFilePerm = 0644
|
||||||
|
)
|
||||||
|
|
||||||
|
var (
|
||||||
|
dataFileFlags = os.O_APPEND | os.O_CREATE | os.O_RDWR
|
||||||
|
indexFileFlags = os.O_APPEND | os.O_CREATE | os.O_RDWR
|
||||||
|
metadataFileFlags = os.O_CREATE | os.O_RDWR
|
||||||
|
|
||||||
|
// errClosed is returned if an operation attempts to manipulate the
|
||||||
|
// database after it has been closed.
|
||||||
|
errClosed = errors.New("vector database already closed")
|
||||||
|
)
|
||||||
|
|
||||||
|
// A VectorDB is a rawData store for storing sequences of binary blobs.
|
||||||
|
//
|
||||||
|
// Items are sequentially added and removed from the VectorDB, but
|
||||||
|
// provides random access to the elements contained within its bounds.
|
||||||
|
type VectorDB struct {
|
||||||
|
// The path the database lives at.
|
||||||
|
path string
|
||||||
|
// Metadata about the database instance
|
||||||
|
metadata *metadata
|
||||||
|
// The number of items stored in the database.
|
||||||
|
items uint64
|
||||||
|
// Mutex protecting the rawData file descriptors
|
||||||
|
lock sync.RWMutex
|
||||||
|
// The file used to rawIndex the content in the rawData file.
|
||||||
|
index *os.File
|
||||||
|
// The file used to store rawData.
|
||||||
|
data *os.File
|
||||||
|
}
|
||||||
|
|
||||||
|
// metadata contains information about the database.
|
||||||
|
type metadata struct {
|
||||||
|
version uint64
|
||||||
|
}
|
||||||
|
|
||||||
|
// indexEntry contains the metadata associated with a stored rawData item.
|
||||||
|
type indexEntry struct {
|
||||||
|
// The position the rawData item starts at within the rawData file.
|
||||||
|
offset uint64
|
||||||
|
// The length of the rawData item in the rawData file.
|
||||||
|
length uint64
|
||||||
|
}
|
||||||
|
|
||||||
|
// indexOffset returns the file offset that corresponse to the blob
|
||||||
|
// located at the specified position in the database.
|
||||||
|
func indexOffset(pos int64) int64 {
|
||||||
|
return pos * indexEntryLen
|
||||||
|
}
|
||||||
|
|
||||||
|
// unmarshallBinary deserializes binary b into the rawIndex entry.
|
||||||
|
func (e *indexEntry) unmarshalBinary(b []byte) error {
|
||||||
|
e.offset = binary.BigEndian.Uint64(b[:8])
|
||||||
|
e.length = binary.BigEndian.Uint64(b[8:16])
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// marshallBinary serializes the rawIndex entry into binary.
|
||||||
|
func (e *indexEntry) marshallBinary() []byte {
|
||||||
|
b := make([]byte, indexEntryLen)
|
||||||
|
binary.BigEndian.PutUint64(b[:8], e.offset)
|
||||||
|
binary.BigEndian.PutUint64(b[8:16], e.length)
|
||||||
|
return b
|
||||||
|
}
|
||||||
|
|
||||||
|
// Open opens a database instance with the specified name
|
||||||
|
// at the provided path.
|
||||||
|
func Open(name, path string) (*VectorDB, error) {
|
||||||
|
databasePath := filepath.Join(path, name)
|
||||||
|
fi, err := os.Stat(databasePath)
|
||||||
|
if os.IsNotExist(err) {
|
||||||
|
if err := os.MkdirAll(databasePath, dbDirPerm); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
} else if err != nil {
|
||||||
|
return nil, err
|
||||||
|
} else if !fi.IsDir() {
|
||||||
|
return nil, fmt.Errorf("open %q: not a directory", databasePath)
|
||||||
|
}
|
||||||
|
|
||||||
|
metadata, err := getOrCreateMetadataFile(path)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
index, err := os.OpenFile(filepath.Join(databasePath, indexFile), indexFileFlags, indexFilePerm)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
data, err := os.OpenFile(filepath.Join(databasePath, dataFile), dataFileFlags, dataFilePerm)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
db := &VectorDB{
|
||||||
|
path: databasePath,
|
||||||
|
metadata: metadata,
|
||||||
|
index: index,
|
||||||
|
data: data,
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := db.repair(); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
return db, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func getOrCreateMetadataFile(path string) (*metadata, error) {
|
||||||
|
metadataFilePath := filepath.Join(path, metadataFile)
|
||||||
|
b, err := ioutil.ReadFile(metadataFilePath)
|
||||||
|
if err == nil {
|
||||||
|
var metadata metadata
|
||||||
|
if err := json.Unmarshal(b, &metadata); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
return &metadata, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
if !os.IsNotExist(err) {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
metadata := metadata{currentVersion}
|
||||||
|
b, err = json.Marshal(metadata)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
ioutil.WriteFile(metadataFilePath, b, metadataFilePerm)
|
||||||
|
return &metadata, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Version returns the current database version.
|
||||||
|
func (db *VectorDB) Version() uint64 {
|
||||||
|
return db.metadata.version
|
||||||
|
}
|
||||||
|
|
||||||
|
// Get retrieves the bytes stored at specified position pos.
|
||||||
|
func (db *VectorDB) Get(pos uint64) ([]byte, error) {
|
||||||
|
db.lock.RLock()
|
||||||
|
defer db.lock.RUnlock()
|
||||||
|
|
||||||
|
if err := db.checkIsOpen(); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := db.checkBounds(pos); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
entry, err := db.indexEntry(pos)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
b := make([]byte, entry.length)
|
||||||
|
if _, err := db.data.ReadAt(b, int64(entry.offset)); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
return b, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *VectorDB) indexEntry(pos uint64) (*indexEntry, error) {
|
||||||
|
b := make([]byte, indexEntryLen)
|
||||||
|
_, err := db.index.ReadAt(b, indexOffset(int64(pos)))
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
entry := new(indexEntry)
|
||||||
|
if err := entry.unmarshalBinary(b); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
return entry, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Append adds the specified blob to the end of the database which should
|
||||||
|
// correspond to the specified pos, which is included as a precaution.
|
||||||
|
//
|
||||||
|
// The result of this operation is not guarranteed to be persisted until
|
||||||
|
// Sync() is called.
|
||||||
|
func (db *VectorDB) Append(pos uint64, blob []byte) error {
|
||||||
|
db.lock.Lock()
|
||||||
|
defer db.lock.Unlock()
|
||||||
|
|
||||||
|
if err := db.checkIsOpen(); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
if pos != db.items {
|
||||||
|
return fmt.Errorf("proposed append position %d does not match current append position %d", pos, db.items)
|
||||||
|
}
|
||||||
|
|
||||||
|
offset, err := db.dataFileSize()
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
if _, err := db.data.Write(blob); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
entry := &indexEntry{uint64(offset), uint64(len(blob))}
|
||||||
|
if _, err := db.index.Write(entry.marshallBinary()); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
db.items++
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *VectorDB) dataFileSize() (int64, error) {
|
||||||
|
fi, err := db.data.Stat()
|
||||||
|
if err != nil {
|
||||||
|
return 0, err
|
||||||
|
}
|
||||||
|
return fi.Size(), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Truncate shortens the database to the desired length items.
|
||||||
|
//
|
||||||
|
// The result of this operation is not guarranteed to be persisted until
|
||||||
|
// Sync() is called.
|
||||||
|
func (db *VectorDB) Truncate(len uint64) error {
|
||||||
|
db.lock.Lock()
|
||||||
|
defer db.lock.Unlock()
|
||||||
|
|
||||||
|
if err := db.checkIsOpen(); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := db.checkBounds(len); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
db.items = len
|
||||||
|
|
||||||
|
newIndexFileSize := len * indexEntryLen
|
||||||
|
if err := db.truncateIndexFile(newIndexFileSize); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
lastEntry, err := db.indexEntry(len - 1)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
newDataFileSize := lastEntry.offset + lastEntry.length
|
||||||
|
if err := db.truncateDataFile(newDataFileSize); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *VectorDB) truncateIndexFile(size uint64) error {
|
||||||
|
return db.index.Truncate(int64(size))
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *VectorDB) truncateDataFile(size uint64) error {
|
||||||
|
return db.data.Truncate(int64(size))
|
||||||
|
}
|
||||||
|
|
||||||
|
// Items returns the length of the database as the number of entries it contains.
|
||||||
|
func (db *VectorDB) Items() uint64 {
|
||||||
|
db.lock.RLock()
|
||||||
|
defer db.lock.RUnlock()
|
||||||
|
|
||||||
|
return db.items
|
||||||
|
}
|
||||||
|
|
||||||
|
// Close closes the database.
|
||||||
|
func (db *VectorDB) Close() error {
|
||||||
|
db.lock.Lock()
|
||||||
|
defer db.lock.Unlock()
|
||||||
|
|
||||||
|
if err := db.checkIsOpen(); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
db.items = 0
|
||||||
|
if err := db.sync(); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
var errs []error
|
||||||
|
if err := db.index.Close(); err != nil {
|
||||||
|
errs = append(errs, fmt.Errorf("error closing rawIndex file: %v", err))
|
||||||
|
}
|
||||||
|
db.index = nil
|
||||||
|
|
||||||
|
if err := db.data.Close(); err != nil {
|
||||||
|
errs = append(errs, fmt.Errorf("error closing rawData file: %v", err))
|
||||||
|
}
|
||||||
|
db.data = nil
|
||||||
|
|
||||||
|
if len(errs) > 0 {
|
||||||
|
return fmt.Errorf("error closing vector database: %v", errs)
|
||||||
|
}
|
||||||
|
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Sync pushes any pending rawData from memory out to disk.
|
||||||
|
//
|
||||||
|
// Note: This is an expensive operation, so use it with care.
|
||||||
|
func (db *VectorDB) Sync() error {
|
||||||
|
db.lock.Lock()
|
||||||
|
defer db.lock.Unlock()
|
||||||
|
|
||||||
|
return db.sync()
|
||||||
|
}
|
||||||
|
|
||||||
|
// sync is the non-thread safe version of Sync.
|
||||||
|
func (db *VectorDB) sync() error {
|
||||||
|
if err := db.checkIsOpen(); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
// Commit rawData before updating indexes.
|
||||||
|
if err := db.data.Sync(); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := db.index.Sync(); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *VectorDB) checkIsOpen() error {
|
||||||
|
if db.index == nil || db.data == nil {
|
||||||
|
return errClosed
|
||||||
|
}
|
||||||
|
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *VectorDB) checkBounds(pos uint64) error {
|
||||||
|
if pos >= db.items {
|
||||||
|
return fmt.Errorf("position out of range (%d >= %d)", pos, db.items)
|
||||||
|
}
|
||||||
|
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *VectorDB) repair() error {
|
||||||
|
db.lock.Lock()
|
||||||
|
defer db.lock.Unlock()
|
||||||
|
|
||||||
|
indexFileSize, err := fileSize(db.index)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
overflow := indexFileSize % indexEntryLen
|
||||||
|
indexFileSize -= overflow
|
||||||
|
if overflow > 0 {
|
||||||
|
if err := db.index.Truncate(indexFileSize); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
dataFileSize, err := fileSize(db.data)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
items := uint64(indexFileSize / indexEntryLen)
|
||||||
|
// Rewind until data file is consistent with what is reported in the index file.
|
||||||
|
for items > 0 {
|
||||||
|
entry, err := db.indexEntry(items - 1)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
// Very likely the index and data files are consistent.
|
||||||
|
if entry.offset+entry.length == uint64(dataFileSize) {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
|
||||||
|
// The index file is ahead of the data file.
|
||||||
|
if entry.offset+entry.length > uint64(dataFileSize) {
|
||||||
|
indexFileSize -= indexEntryLen
|
||||||
|
if err := db.index.Truncate(indexFileSize); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
items--
|
||||||
|
break
|
||||||
|
}
|
||||||
|
|
||||||
|
// The last blob in the data file must be corrupt.
|
||||||
|
dataFileSize = int64(entry.offset + entry.length)
|
||||||
|
if err := db.data.Truncate(dataFileSize); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
db.items = items
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func fileSize(file *os.File) (int64, error) {
|
||||||
|
fi, err := file.Stat()
|
||||||
|
if err != nil {
|
||||||
|
return -1, err
|
||||||
|
}
|
||||||
|
|
||||||
|
return fi.Size(), nil
|
||||||
|
}
|
||||||
414
ethdb/vectordb/vectordb_test.go
Normal file
414
ethdb/vectordb/vectordb_test.go
Normal file
|
|
@ -0,0 +1,414 @@
|
||||||
|
// Copyright 2019 The go-ethereum Authors
|
||||||
|
// This file is part of the go-ethereum library.
|
||||||
|
//
|
||||||
|
// The go-ethereum library is free software: you can redistribute it and/or modify
|
||||||
|
// it under the terms of the GNU Lesser General Public License as published by
|
||||||
|
// the Free Software Foundation, either version 3 of the License, or
|
||||||
|
// (at your option) any later version.
|
||||||
|
//
|
||||||
|
// The go-ethereum library 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 Lesser General Public License for more details.
|
||||||
|
//
|
||||||
|
// You should have received a copy of the GNU Lesser General Public License
|
||||||
|
// along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.
|
||||||
|
package vectordb
|
||||||
|
|
||||||
|
import (
|
||||||
|
"encoding/hex"
|
||||||
|
"io/ioutil"
|
||||||
|
"os"
|
||||||
|
"path/filepath"
|
||||||
|
"testing"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestOpenVersion(t *testing.T) {
|
||||||
|
dir, rmdir := createTempDir(t)
|
||||||
|
defer rmdir()
|
||||||
|
|
||||||
|
vectorDB, err := Open("vectordb", dir)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Open(%q) = %v, want <nil>", dir, err)
|
||||||
|
}
|
||||||
|
defer vectorDB.Close()
|
||||||
|
|
||||||
|
if version := vectorDB.Version(); version != currentVersion {
|
||||||
|
t.Fatalf("vectorDB.Version() = %d, want %d", version, currentVersion)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestVectorDB_Repair(t *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
rawIndex, rawData func() []byte
|
||||||
|
blobs [][]byte
|
||||||
|
}{
|
||||||
|
{
|
||||||
|
"LastIndexEntryMissingAByte",
|
||||||
|
func() []byte {
|
||||||
|
index := marshalIndexEntries(&indexEntry{0, 4}, &indexEntry{4, 3})
|
||||||
|
return index[:len(index)-1]
|
||||||
|
},
|
||||||
|
func() []byte {
|
||||||
|
return []byte{1, 1, 1, 1, 2, 2, 2}
|
||||||
|
},
|
||||||
|
[][]byte{{1, 1, 1, 1}},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"DanglingIndexEntry",
|
||||||
|
func() []byte {
|
||||||
|
index := marshalIndexEntries(&indexEntry{0, 4}, &indexEntry{4, 3})
|
||||||
|
return index[:len(index)-1]
|
||||||
|
},
|
||||||
|
func() []byte {
|
||||||
|
return []byte{1, 1, 1, 1}
|
||||||
|
},
|
||||||
|
[][]byte{{1, 1, 1, 1}},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"LastDataEntryMissingAByte",
|
||||||
|
func() []byte {
|
||||||
|
return marshalIndexEntries(&indexEntry{0, 4}, &indexEntry{4, 3})
|
||||||
|
},
|
||||||
|
func() []byte {
|
||||||
|
return []byte{1, 1, 1, 1, 2, 2}
|
||||||
|
},
|
||||||
|
[][]byte{{1, 1, 1, 1}},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tc := range tests {
|
||||||
|
t.Run(tc.name, func(t *testing.T) {
|
||||||
|
dir, rmdir := createTempDir(t)
|
||||||
|
defer rmdir()
|
||||||
|
|
||||||
|
if err := os.MkdirAll(filepath.Join(dir, "vectordb"), dbDirPerm); err != nil {
|
||||||
|
t.Fatalf("Error creating mock database directory: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := ioutil.WriteFile(filepath.Join(dir, "vectordb", indexFile), tc.rawIndex(), indexFilePerm); err != nil {
|
||||||
|
t.Fatalf("Error writing mock rawIndex file: %v", err)
|
||||||
|
}
|
||||||
|
if err := ioutil.WriteFile(filepath.Join(dir, "vectordb", dataFile), tc.rawData(), dataFilePerm); err != nil {
|
||||||
|
t.Fatalf("Error writing mock rawData file: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
vectorDB, err := Open("vectordb", dir)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Open(%q) = %v, got <nil>", dir, err)
|
||||||
|
}
|
||||||
|
defer vectorDB.Close()
|
||||||
|
|
||||||
|
if got := vectorDB.Items(); got != uint64(len(tc.blobs)) {
|
||||||
|
t.Fatalf("vectorDB.Items() = %d, want %d", got, uint64(len(tc.blobs)))
|
||||||
|
}
|
||||||
|
|
||||||
|
for i, want := range tc.blobs {
|
||||||
|
got, err := vectorDB.Get(uint64(i))
|
||||||
|
if err != nil {
|
||||||
|
t.Errorf("vectorDB.Get(%d) = %s, %v, want %s, <nil>", uint64(i), hex.EncodeToString(got), err, hex.EncodeToString(want))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func marshalIndexEntries(entries ...*indexEntry) []byte {
|
||||||
|
var b []byte
|
||||||
|
for _, entry := range entries {
|
||||||
|
b = append(b, entry.marshallBinary()...)
|
||||||
|
}
|
||||||
|
return b
|
||||||
|
}
|
||||||
|
|
||||||
|
func marshalDataBlobs(blobs ...[]byte) []byte {
|
||||||
|
var b []byte
|
||||||
|
for _, blob := range blobs {
|
||||||
|
b = append(b, blob...)
|
||||||
|
}
|
||||||
|
return b
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestOpen_DirectoryAlreadyExists_ReturnsError(t *testing.T) {
|
||||||
|
dir, rmdir := createTempDir(t)
|
||||||
|
defer rmdir()
|
||||||
|
|
||||||
|
vectorDB, err := Open("vectordb", dir)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Open(%q) = %v, want <nil>", dir, err)
|
||||||
|
}
|
||||||
|
vectorDB.Close()
|
||||||
|
|
||||||
|
vectorDB2, err := Open("vectordb", dir)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Open(%q) = %v, want <nil>", dir, err)
|
||||||
|
}
|
||||||
|
vectorDB2.Close()
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestVectorDB_AppendGet(t *testing.T) {
|
||||||
|
blobs := [][]byte{
|
||||||
|
{1},
|
||||||
|
{2, 2},
|
||||||
|
{3, 3, 3},
|
||||||
|
}
|
||||||
|
|
||||||
|
dir, rmdir := createTempDir(t)
|
||||||
|
defer rmdir()
|
||||||
|
|
||||||
|
vectorDB, err := Open("vectordb", dir)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Open(%q) = %v, want <nil>", dir, err)
|
||||||
|
}
|
||||||
|
defer vectorDB.Close()
|
||||||
|
|
||||||
|
for i, blob := range blobs {
|
||||||
|
if err := vectorDB.Append(uint64(i), blob); err != nil {
|
||||||
|
t.Errorf("vectorDB.Append(%d, %q) = %v, want <nil>", i, hex.EncodeToString(blob), err)
|
||||||
|
}
|
||||||
|
if vectorDB.Items() != uint64(i+1) {
|
||||||
|
t.Errorf("vectorDB.Items() = %d, want %d", vectorDB.Items(), uint64(i+1))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
for i, want := range blobs {
|
||||||
|
got, err := vectorDB.Get(uint64(i))
|
||||||
|
if err != nil {
|
||||||
|
t.Errorf("vectorDB.Get(%d) = %s, %v, want %s, <nil>", uint64(i), hex.EncodeToString(got), err, hex.EncodeToString(want))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestVectorDB_GetOnExistingDatabase(t *testing.T) {
|
||||||
|
blobs := [][]byte{
|
||||||
|
{1},
|
||||||
|
{2, 2},
|
||||||
|
{3, 3, 3},
|
||||||
|
}
|
||||||
|
|
||||||
|
dir, rmdir := createTempDir(t)
|
||||||
|
defer rmdir()
|
||||||
|
|
||||||
|
vectorDB, err := Open("vectordb", dir)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Open(%q) = %v, want <nil>", dir, err)
|
||||||
|
}
|
||||||
|
|
||||||
|
for i, blob := range blobs {
|
||||||
|
if err := vectorDB.Append(uint64(i), blob); err != nil {
|
||||||
|
t.Errorf("vectorDB.Append(%d, %q) = %v, want <nil>", i, hex.EncodeToString(blob), err)
|
||||||
|
}
|
||||||
|
if vectorDB.Items() != uint64(i+1) {
|
||||||
|
t.Errorf("vectorDB.Items() = %d, want %d", vectorDB.Items(), uint64(i+1))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := vectorDB.Sync(); err != nil {
|
||||||
|
t.Fatalf("vectorDB.Sync() = %v, want <nil>", err)
|
||||||
|
}
|
||||||
|
vectorDB.Close()
|
||||||
|
|
||||||
|
vectorDB2, err := Open("vectordb", dir)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Open(%q) = %v, want <nil>", dir, err)
|
||||||
|
}
|
||||||
|
defer vectorDB2.Close()
|
||||||
|
|
||||||
|
for i, want := range blobs {
|
||||||
|
got, err := vectorDB2.Get(uint64(i))
|
||||||
|
if err != nil {
|
||||||
|
t.Errorf("vectorDB2.Get(%d) = %s, %v, want %s, <nil>", uint64(i), hex.EncodeToString(got), err, hex.EncodeToString(want))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAppend_PositionMismatch_ReturnsError(t *testing.T) {
|
||||||
|
items := [][]byte{{1}, {2}, {3}}
|
||||||
|
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
pos uint64
|
||||||
|
}{
|
||||||
|
{
|
||||||
|
"Before",
|
||||||
|
uint64(len(items) - 1),
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"After",
|
||||||
|
uint64(len(items) + 1),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tc := range tests {
|
||||||
|
t.Run(tc.name, func(t *testing.T) {
|
||||||
|
dir, rmdir := createTempDir(t)
|
||||||
|
defer rmdir()
|
||||||
|
|
||||||
|
vectorDB, err := Open("vectordb", dir)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Open(%q) = %v, want <nil>", dir, err)
|
||||||
|
}
|
||||||
|
defer vectorDB.Close()
|
||||||
|
|
||||||
|
for i, item := range items {
|
||||||
|
if err := vectorDB.Append(uint64(i), item); err != nil {
|
||||||
|
t.Errorf("vectorDB.Append(%d, %q) = %v, want <nil>", i, hex.EncodeToString(item), err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := vectorDB.Append(tc.pos, []byte{0}); err == nil {
|
||||||
|
t.Fatalf("vector.Append(%d, %s) = %v, want <error>", tc.pos, hex.EncodeToString([]byte{0}), err)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestVectorDB_GetGreaterThanLen_ReturnsError(t *testing.T) {
|
||||||
|
dir, rmdir := createTempDir(t)
|
||||||
|
defer rmdir()
|
||||||
|
|
||||||
|
vectorDB, err := Open("vectordb", dir)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Open(%q) = %v, want <nil>", dir, err)
|
||||||
|
}
|
||||||
|
defer vectorDB.Close()
|
||||||
|
|
||||||
|
for i := 0; i < 3; i++ {
|
||||||
|
vectorDB.Append(uint64(i), []byte{1, 2, 3})
|
||||||
|
}
|
||||||
|
|
||||||
|
if got, err := vectorDB.Get(3); err == nil {
|
||||||
|
t.Errorf("vectorDB.Get(%d) = %s, %v, want \"\", <err>", uint64(3), hex.EncodeToString(got), err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestVectorDB_Truncate(t *testing.T) {
|
||||||
|
const truncatedLen = 2
|
||||||
|
blobs := [][]byte{
|
||||||
|
{1},
|
||||||
|
{2, 2},
|
||||||
|
{3, 3, 3},
|
||||||
|
}
|
||||||
|
|
||||||
|
dir, rmdir := createTempDir(t)
|
||||||
|
defer rmdir()
|
||||||
|
|
||||||
|
vectorDB, err := Open("vectordb", dir)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Open(%q) = %v, want <nil>", dir, err)
|
||||||
|
}
|
||||||
|
defer vectorDB.Close()
|
||||||
|
|
||||||
|
for i, blob := range blobs {
|
||||||
|
if err := vectorDB.Append(uint64(i), blob); err != nil {
|
||||||
|
t.Fatalf("vectorDB.Append(%d, %q) = %v, want <nil>", i, hex.EncodeToString(blob), err)
|
||||||
|
}
|
||||||
|
if vectorDB.Items() != uint64(i+1) {
|
||||||
|
t.Fatalf("vectorDB.Items() = %d, want %d", vectorDB.Items(), uint64(i+1))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := vectorDB.Truncate(truncatedLen); err != nil {
|
||||||
|
t.Fatalf("vectorDB.Truncate(%d) = %v, want <nil>", truncatedLen, err)
|
||||||
|
}
|
||||||
|
|
||||||
|
for i, want := range blobs[:truncatedLen] {
|
||||||
|
got, err := vectorDB.Get(uint64(i))
|
||||||
|
if err != nil {
|
||||||
|
t.Errorf("vectorDB.Get(%d) = %s, %v, want %s, <nil>", uint64(i), hex.EncodeToString(got), err, hex.EncodeToString(want))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if got, err := vectorDB.Get(truncatedLen); err == nil {
|
||||||
|
t.Errorf("vectorDB.Get(%d) = %s, %v, want \"\", <err>", truncatedLen, hex.EncodeToString(got), err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestVectorDB_TruncateGreaterThanLen_ReturnsError(t *testing.T) {
|
||||||
|
dir, rmdir := createTempDir(t)
|
||||||
|
defer rmdir()
|
||||||
|
|
||||||
|
vectorDB, err := Open("vectordb", dir)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Open(%q) = %v, want <nil>", dir, err)
|
||||||
|
}
|
||||||
|
defer vectorDB.Close()
|
||||||
|
|
||||||
|
for i := 0; i < 3; i++ {
|
||||||
|
vectorDB.Append(uint64(i), []byte{1, 2, 3})
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := vectorDB.Truncate(3); err == nil {
|
||||||
|
t.Errorf("vectorDB.Truncate(%d) = %v, want <err>", uint64(3), err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestVectorDB_ReturnsErrrWhenClosed(t *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
do func(*VectorDB) error
|
||||||
|
}{
|
||||||
|
{
|
||||||
|
"Append",
|
||||||
|
func(db *VectorDB) error {
|
||||||
|
return db.Append(uint64(0), []byte{1, 2, 3})
|
||||||
|
},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"Get",
|
||||||
|
func(db *VectorDB) error {
|
||||||
|
_, err := db.Get(0)
|
||||||
|
return err
|
||||||
|
},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"Truncate",
|
||||||
|
func(db *VectorDB) error {
|
||||||
|
return db.Truncate(0)
|
||||||
|
},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"Sync",
|
||||||
|
func(db *VectorDB) error {
|
||||||
|
return db.Sync()
|
||||||
|
},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"Close",
|
||||||
|
func(db *VectorDB) error {
|
||||||
|
return db.Close()
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tc := range tests {
|
||||||
|
t.Run(tc.name, func(t *testing.T) {
|
||||||
|
dir, rmdir := createTempDir(t)
|
||||||
|
defer rmdir()
|
||||||
|
|
||||||
|
vectorDB, err := Open("vectordb", dir)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Open(%q) = %v, want <nil>", dir, err)
|
||||||
|
}
|
||||||
|
|
||||||
|
vectorDB.Close()
|
||||||
|
|
||||||
|
if err := tc.do(vectorDB); err != errClosed {
|
||||||
|
t.Fatalf("vectorDB.%s = %v, want %v", tc.name, err, errClosed)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func createTempDir(t *testing.T) (string, func()) {
|
||||||
|
t.Helper()
|
||||||
|
|
||||||
|
root, err := ioutil.TempDir(os.TempDir(), "vectordb_test_")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Error creating test directory: %v", err)
|
||||||
|
}
|
||||||
|
return root, func() {
|
||||||
|
os.RemoveAll(root)
|
||||||
|
}
|
||||||
|
}
|
||||||
Loading…
Reference in a new issue