Skip to content

Commit b5a9595

Browse files
krasunkrasun
committed
Refactor disk table search to use io.Reader
1 parent 8ea80a5 commit b5a9595

File tree

2 files changed

+65
-33
lines changed

2 files changed

+65
-33
lines changed

‎disktable.go

+42-30
Original file line numberDiff line numberDiff line change
@@ -73,7 +73,13 @@ func searchInDiskTable(dbDir string, index int, key []byte) ([]byte, bool, error
7373
prefix := strconv.Itoa(index) + "-"
7474

7575
sparseIndexPath := path.Join(dbDir, prefix+diskTableSparseIndexFileName)
76-
from, to, ok, err := searchInSparseIndex(sparseIndexPath, key)
76+
sparseIndexFile, err := os.OpenFile(sparseIndexPath, os.O_RDONLY, 0600)
77+
if err != nil {
78+
return nil, false, fmt.Errorf("failed to open sparse index file: %w", err)
79+
}
80+
defer sparseIndexFile.Close()
81+
82+
from, to, ok, err := searchInSparseIndex(sparseIndexFile, key)
7783
if err != nil {
7884
return nil, false, fmt.Errorf("failed to search in sparse index file %s: %w", sparseIndexPath, err)
7985
}
@@ -82,7 +88,13 @@ func searchInDiskTable(dbDir string, index int, key []byte) ([]byte, bool, error
8288
}
8389

8490
indexPath := path.Join(dbDir, prefix+diskTableIndexFileName)
85-
offset, ok, err := searchInIndex(indexPath, from, to, key)
91+
indexFile, err := os.OpenFile(indexPath, os.O_RDONLY, 0600)
92+
if err != nil {
93+
return nil, false, fmt.Errorf("failed to open index file: %w", err)
94+
}
95+
defer indexFile.Close()
96+
97+
offset, ok, err := searchInIndex(indexFile, from, to, key)
8698
if err != nil {
8799
return nil, false, fmt.Errorf("failed to search in index file %s: %w", indexPath, err)
88100
}
@@ -91,29 +103,41 @@ func searchInDiskTable(dbDir string, index int, key []byte) ([]byte, bool, error
91103
}
92104

93105
dataPath := path.Join(dbDir, prefix+diskTableDataFileName)
94-
value, ok, err := searchInDataFile(dataPath, offset, key)
106+
dataFile, err := os.OpenFile(dataPath, os.O_RDONLY, 0600)
107+
if err != nil {
108+
return nil, false, fmt.Errorf("failed to open data file: %w", err)
109+
}
110+
defer dataFile.Close()
111+
112+
value, ok, err := searchInDataFile(dataFile, offset, key)
95113
if err != nil {
96114
return nil, false, fmt.Errorf("failed to search in data file %s: %w", dataPath, err)
97115
}
98116

117+
if err := sparseIndexFile.Close(); err != nil {
118+
return nil, false, fmt.Errorf("failed to close sparse index file: %w", err)
119+
}
120+
121+
if err := indexFile.Close(); err != nil {
122+
return nil, false, fmt.Errorf("failed to close index file: %w", err)
123+
}
124+
125+
if err := dataFile.Close(); err != nil {
126+
return nil, false, fmt.Errorf("failed to close data file: %w", err)
127+
}
128+
99129
return value, ok, nil
100130
}
101131

102132
// searchInDataFile searches a value by the key in the data file from the given offset.
103133
// The offset must always point to the beginning of the record.
104-
func searchInDataFile(path string, offset int, searchKey []byte) ([]byte, bool, error) {
105-
f, err := os.OpenFile(path, os.O_RDONLY, 0600)
106-
if err != nil {
107-
return nil, false, err
108-
}
109-
defer f.Close()
110-
111-
if _, err := f.Seek(int64(offset), io.SeekCurrent); err != nil {
134+
func searchInDataFile(r io.ReadSeeker, offset int, searchKey []byte) ([]byte, bool, error) {
135+
if _, err := r.Seek(int64(offset), io.SeekStart); err != nil {
112136
return nil, false, fmt.Errorf("failed to seek: %w", err)
113137
}
114138

115139
for {
116-
key, value, err := decode(f)
140+
key, value, err := decode(r)
117141
if err != nil && err != io.EOF {
118142
return nil, false, fmt.Errorf("failed to read: %w", err)
119143
}
@@ -128,19 +152,13 @@ func searchInDataFile(path string, offset int, searchKey []byte) ([]byte, bool,
128152
}
129153

130154
// searchInIndex searches key in the index file in specified range.
131-
func searchInIndex(path string, from, to int, searchKey []byte) (int, bool, error) {
132-
f, err := os.OpenFile(path, os.O_RDONLY, 0600)
133-
if err != nil {
134-
return 0, false, err
135-
}
136-
defer f.Close()
137-
138-
if _, err := f.Seek(int64(from), io.SeekStart); err != nil {
155+
func searchInIndex(r io.ReadSeeker, from, to int, searchKey []byte) (int, bool, error) {
156+
if _, err := r.Seek(int64(from), io.SeekStart); err != nil {
139157
return 0, false, fmt.Errorf("failed to seek: %w", err)
140158
}
141159

142160
for {
143-
key, value, err := decode(f)
161+
key, value, err := decode(r)
144162
if err != nil && err != io.EOF {
145163
return 0, false, fmt.Errorf("failed to read: %w", err)
146164
}
@@ -154,7 +172,7 @@ func searchInIndex(path string, from, to int, searchKey []byte) (int, bool, erro
154172
}
155173

156174
if to > from {
157-
current, err := f.Seek(0, io.SeekCurrent)
175+
current, err := r.Seek(0, io.SeekCurrent)
158176
if err != nil {
159177
return 0, false, fmt.Errorf("failed to seek: %w", err)
160178
}
@@ -167,16 +185,10 @@ func searchInIndex(path string, from, to int, searchKey []byte) (int, bool, erro
167185
}
168186

169187
// searchInSparseIndex searches a range between which the key is located.
170-
func searchInSparseIndex(path string, searchKey []byte) (int, int, bool, error) {
171-
f, err := os.OpenFile(path, os.O_RDONLY, 0600)
172-
if err != nil {
173-
return 0, 0, false, err
174-
}
175-
defer f.Close()
176-
188+
func searchInSparseIndex(r io.Reader, searchKey []byte) (int, int, bool, error) {
177189
from := -1
178190
for {
179-
key, value, err := decode(f)
191+
key, value, err := decode(r)
180192
if err != nil && err != io.EOF {
181193
return 0, 0, false, fmt.Errorf("failed to read: %w", err)
182194
}

‎disktable_test.go

+23-3
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package lsmtree
33
import (
44
"bytes"
55
"fmt"
6+
"io"
67
"io/ioutil"
78
"os"
89
"path"
@@ -113,8 +114,14 @@ func TestSearchInDataFile(t *testing.T) {
113114
{[]byte("k"), nil, false, false, 0},
114115
}
115116

117+
dataFile, err := os.OpenFile(path.Join(dbDir, "0-data.db"), os.O_RDONLY, 0600)
118+
if err != nil {
119+
t.Fatalf("failed to open data file: %s", err)
120+
}
121+
defer dataFile.Close()
122+
116123
for _, c := range cases {
117-
value, ok, err := searchInDataFile(path.Join(dbDir, "0-data.db"), c.offset, c.key)
124+
value, ok, err := searchInDataFile(dataFile, c.offset, c.key)
118125
if !((c.value == nil && value == nil) || (bytes.Equal(c.value, value))) {
119126
t.Fatalf("values do not match for %s, err = %v: %s != %s", string(c.key), err, string(c.value), string(value))
120127
}
@@ -148,8 +155,14 @@ func TestSearchInIndex(t *testing.T) {
148155
{[]byte("k"), 150, 0, false, false, 0},
149156
}
150157

158+
indexFile, err := os.OpenFile(path.Join(dbDir, "0-index.db"), os.O_RDONLY, 0600)
159+
if err != nil {
160+
t.Fatalf("failed to open index file: %s", err)
161+
}
162+
defer indexFile.Close()
163+
151164
for _, c := range cases {
152-
offset, ok, err := searchInIndex(path.Join(dbDir, "0-index.db"), c.from, c.to, c.key)
165+
offset, ok, err := searchInIndex(indexFile, c.from, c.to, c.key)
153166
if c.offset != offset {
154167
t.Fatalf("offset does not match for %s, err = %v: %d != %d", string(c.key), err, c.offset, offset)
155168
}
@@ -182,8 +195,15 @@ func TestSearchInSparseIndex(t *testing.T) {
182195
{[]byte("k"), 150, 0, true, false},
183196
}
184197

198+
sparseIndexFile, err := os.OpenFile(path.Join(dbDir, "0-sparse.db"), os.O_RDONLY, 0600)
199+
if err != nil {
200+
t.Fatalf("failed to open sparse index file: %s", err)
201+
}
202+
defer sparseIndexFile.Close()
203+
185204
for _, c := range cases {
186-
from, to, ok, err := searchInSparseIndex(path.Join(dbDir, "0-sparse.db"), c.key)
205+
sparseIndexFile.Seek(0, io.SeekStart)
206+
from, to, ok, err := searchInSparseIndex(sparseIndexFile, c.key)
187207
if c.from != from || c.to != to {
188208
t.Fatalf("from and to do not match for %s, err = %v: %d != %d or %d != %d", string(c.key), err, c.from, from, c.to, to)
189209
}

0 commit comments

Comments
 (0)