diff --git a/ethdb/vectordb/vectordb.go b/ethdb/vectordb/vectordb.go new file mode 100644 index 0000000000..4c12c5008c --- /dev/null +++ b/ethdb/vectordb/vectordb.go @@ -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 . + +// 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 +} diff --git a/ethdb/vectordb/vectordb_test.go b/ethdb/vectordb/vectordb_test.go new file mode 100644 index 0000000000..e1f4ce0180 --- /dev/null +++ b/ethdb/vectordb/vectordb_test.go @@ -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 . +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 ", 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 ", 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, ", 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 ", dir, err) + } + vectorDB.Close() + + vectorDB2, err := Open("vectordb", dir) + if err != nil { + t.Fatalf("Open(%q) = %v, want ", 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 ", 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 ", 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, ", 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 ", dir, err) + } + + for i, blob := range blobs { + if err := vectorDB.Append(uint64(i), blob); err != nil { + t.Errorf("vectorDB.Append(%d, %q) = %v, want ", 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 ", err) + } + vectorDB.Close() + + vectorDB2, err := Open("vectordb", dir) + if err != nil { + t.Fatalf("Open(%q) = %v, want ", 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, ", 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 ", 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 ", i, hex.EncodeToString(item), err) + } + } + + if err := vectorDB.Append(tc.pos, []byte{0}); err == nil { + t.Fatalf("vector.Append(%d, %s) = %v, want ", 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 ", 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 \"\", ", 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 ", 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 ", 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 ", 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, ", 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 \"\", ", 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 ", 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 ", 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 ", 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) + } +}