Skip to content

Commit f204338

Browse files
reflogIlya Yakelzon
andauthored
perf: pre-size fileSource reads (close the InitFrom io.ReadAll churn) (#10)
* perf: pre-size fileSource reads too (close the InitFrom churn) Follow-up to #9. That PR made webSource and tarGzSource size-aware, but fileSource was left returning a bare *os.File, so readAll's knownSize() returned false for file reads and fell back to io.ReadAll's doubling growth. geo loads its cached ~75MB MaxMind mmdb from disk via Runner.InitFrom(FromFile(path)) at startup, so that read churned through the 75->60->48->38->30->24MB realloc staircase on every boot (confirmed via alloc_space pprof on a lantern-box v0.0.91 box: io.ReadAll 373MB under keepcurrent.readAll <- syncOnce <- InitFrom <- geo.FromWeb). Make fileSource.Fetch return a sizedReadCloser carrying the stat size when there's no preprocessor, so the read is a single pre-sized allocation. The preprocessor path keeps the io.ReadAll fallback (a preprocessor can change the byte count). Also close the file on the early-return/error paths, which previously leaked the fd. go build/vet/test green; adds TestFileSourceIsSizeAware. * review: fix inverted fileSource modified-since check, close-couple preprocessor stream, hermetic web test - fileSource.Fetch returned ErrUnmodified exactly when the file WAS newer than the cutoff (and fetched when unchanged). Invert to match the Source.Fetch contract and byteSource.Fetch. - couple a preprocessor's returned stream's Close to the underlying *os.File so io.NopCloser-style preprocessors don't leak the descriptor. - add TestFileSourceUnmodified covering the ErrUnmodified path. - make TestUpdateFromWeb hermetic via httptest; the live httpbin.org dependency intermittently hung and timed out CI. --------- Co-authored-by: Ilya Yakelzon <reflog@getlantern.org>
1 parent 0364abd commit f204338

3 files changed

Lines changed: 93 additions & 7 deletions

File tree

‎keepcurrent_test.go‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,11 @@ import (
66
"compress/gzip"
77
"crypto/rand"
88
"encoding/json"
9+
"fmt"
910
"io"
1011
"io/ioutil"
12+
"net/http"
13+
"net/http/httptest"
1114
"os"
1215
"strings"
1316
"sync/atomic"
@@ -103,7 +106,13 @@ func writeTempFile(t *testing.T, b []byte) (string, []byte) {
103106

104107
func TestUpdateFromWeb(t *testing.T) {
105108
ch := make(chan []byte)
106-
url := "https://httpbin.org/get"
109+
// Serve an httpbin-/get-shaped body locally so the test stays hermetic —
110+
// the live httpbin.org dependency intermittently hung and timed out CI.
111+
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
112+
fmt.Fprintf(w, `{"url": %q}`, "http://"+r.Host+r.URL.String())
113+
}))
114+
defer srv.Close()
115+
url := srv.URL + "/get"
107116
runner := New(FromWeb(url), ToChannel(ch))
108117
runner.OnSourceError = func(err error, tries int) time.Duration {
109118
assert.Fail(t, "unexpected source error "+err.Error())

‎read_test.go‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,12 +3,59 @@ package keepcurrent
33
import (
44
"bytes"
55
"io"
6+
"os"
7+
"path/filepath"
68
"testing"
9+
"time"
710

811
"github.com/stretchr/testify/assert"
912
"github.com/stretchr/testify/require"
1013
)
1114

15+
// FromFile (no preprocessor) must hand back a size-aware reader so the Runner's
16+
// readAll pre-sizes its buffer rather than falling back to io.ReadAll's
17+
// doubling growth — the cached-mmdb-at-startup case that regressed otherwise.
18+
func TestFileSourceIsSizeAware(t *testing.T) {
19+
payload := bytes.Repeat([]byte("m"), 3*1024*1024) // 3 MiB
20+
path := filepath.Join(t.TempDir(), "cached.bin")
21+
require.NoError(t, os.WriteFile(path, payload, 0o644))
22+
23+
rc, err := FromFile(path).Fetch(time.Time{})
24+
require.NoError(t, err)
25+
defer rc.Close()
26+
27+
n, ok := knownSize(rc)
28+
assert.True(t, ok, "FromFile reader should report its size")
29+
assert.EqualValues(t, len(payload), n)
30+
31+
got, err := readAll(rc)
32+
require.NoError(t, err)
33+
assert.Equal(t, payload, got)
34+
assert.LessOrEqual(t, cap(got), len(payload)+bytes.MinRead, "must be a single pre-sized allocation, not grown")
35+
}
36+
37+
// Fetch must honour its "only if modified since" contract: when ifNewerThan is
38+
// at or after the file's modtime, it returns ErrUnmodified and leaves no reader
39+
// (hence no leaked descriptor) for the caller to close.
40+
func TestFileSourceUnmodified(t *testing.T) {
41+
path := filepath.Join(t.TempDir(), "cached.bin")
42+
require.NoError(t, os.WriteFile(path, []byte("payload"), 0o644))
43+
44+
fi, err := os.Stat(path)
45+
require.NoError(t, err)
46+
47+
// A cutoff strictly after the modtime => unmodified.
48+
rc, err := FromFile(path).Fetch(fi.ModTime().Add(time.Hour))
49+
assert.ErrorIs(t, err, ErrUnmodified)
50+
assert.Nil(t, rc)
51+
52+
// A cutoff strictly before the modtime => the file is newer, so we fetch.
53+
rc, err = FromFile(path).Fetch(fi.ModTime().Add(-time.Hour))
54+
require.NoError(t, err)
55+
require.NotNil(t, rc)
56+
require.NoError(t, rc.Close())
57+
}
58+
1259
// unsizedReader hides any size its wrapped reader might otherwise expose (it
1360
// implements only Read), forcing readAll down the io.ReadAll fallback path.
1461
type unsizedReader struct{ r io.Reader }

‎source.go‎

Lines changed: 36 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -157,6 +157,22 @@ type fileSource struct {
157157
preprocessor func(io.ReadCloser) (io.ReadCloser, error)
158158
}
159159

160+
// fileBackedReadCloser ties a preprocessed stream's lifetime to the file it was
161+
// derived from, closing both when Close is called. A double close of the file
162+
// (when the preprocessor's stream already closes it) is tolerated.
163+
type fileBackedReadCloser struct {
164+
io.ReadCloser
165+
f *os.File
166+
}
167+
168+
func (r fileBackedReadCloser) Close() error {
169+
err := r.ReadCloser.Close()
170+
if cerr := r.f.Close(); cerr != nil && !errors.Is(cerr, os.ErrClosed) && err == nil {
171+
err = cerr
172+
}
173+
return err
174+
}
175+
160176
// FromFile constructs a source from the given file path.
161177
func FromFile(path string) Source {
162178
return &fileSource{path, nil}
@@ -174,18 +190,32 @@ func (s *fileSource) Fetch(ifNewerThan time.Time) (io.ReadCloser, error) {
174190
}
175191
fi, err := f.Stat()
176192
if err != nil {
193+
f.Close()
177194
return nil, err
178195
}
179-
if !ifNewerThan.IsZero() && ifNewerThan.Before(fi.ModTime()) {
196+
// Fetch's contract is "fetch only if modified since ifNewerThan", so we
197+
// signal ErrUnmodified when the file is no newer than the cutoff — mirroring
198+
// byteSource.Fetch in the tests. (Earlier this comparison was inverted.)
199+
if !ifNewerThan.IsZero() && !fi.ModTime().After(ifNewerThan) {
200+
f.Close()
180201
return nil, ErrUnmodified
181202
}
182-
var result io.ReadCloser
183-
result = f
184203
if s.preprocessor != nil {
185-
result, err = s.preprocessor(f)
204+
// A preprocessor may change the byte count, so we can't declare a size
205+
// up front; readAll falls back to io.ReadAll for the preprocessed stream.
206+
result, err := s.preprocessor(f)
186207
if err != nil {
208+
f.Close()
187209
return nil, err
188210
}
189-
}
190-
return result, nil
211+
// The preprocessor may hand back a stream that doesn't close the file it
212+
// wraps (e.g. io.NopCloser), so couple Close to the underlying file to
213+
// avoid leaking the descriptor.
214+
return fileBackedReadCloser{ReadCloser: result, f: f}, nil
215+
}
216+
// Surface the file size so the Runner's readAll pre-sizes its buffer instead
217+
// of growing by reallocation. The cached payloads read here are exactly the
218+
// large ones that churn — e.g. geo loads a ~75MB MaxMind mmdb from disk via
219+
// InitFrom(FromFile(...)) at startup.
220+
return sizedReadCloser{ReadCloser: f, n: fi.Size()}, nil
191221
}

0 commit comments

Comments
 (0)