From 5072219829ceb1244dff0bdff56ee00e49ba54a9 Mon Sep 17 00:00:00 2001 From: "mudler's LocalAI [bot]" <139863280+localai-bot@users.noreply.github.com> Date: Fri, 21 Aug 2026 16:14:26 +0200 Subject: [PATCH] feat(xio): make copy buffer size configurable (#11660) * docs: design configurable copy buffering Document the context-aware copy buffer option and its validation plan. Assisted-by: Codex:gpt-5 * feat(xio): configure context copy buffer size --------- Co-authored-by: Ettore Di Giacinto --- ...6-08-21-configurable-copy-buffer-design.md | 73 +++++++++++ pkg/xio/copy.go | 28 ++++- pkg/xio/copy_suite_test.go | 13 ++ pkg/xio/copy_test.go | 115 ++++++++++++++++++ 4 files changed, 226 insertions(+), 3 deletions(-) create mode 100644 docs/superpowers/specs/2026-08-21-configurable-copy-buffer-design.md create mode 100644 pkg/xio/copy_suite_test.go create mode 100644 pkg/xio/copy_test.go diff --git a/docs/superpowers/specs/2026-08-21-configurable-copy-buffer-design.md b/docs/superpowers/specs/2026-08-21-configurable-copy-buffer-design.md new file mode 100644 index 000000000..1085af7fe --- /dev/null +++ b/docs/superpowers/specs/2026-08-21-configurable-copy-buffer-design.md @@ -0,0 +1,73 @@ +# Configurable copy buffer design + +**Date:** 21 August 2026 +**Status:** Approved + +## Problem + +`pkg/xio.Copy` wraps a source reader so a context can stop a copy between +reads. It delegates to `io.Copy`, which uses a 32 KiB buffer for the wrapped +reader and writer types used by model downloads. + +Small writes limit model import throughput when the models directory uses an +SMB volume. The development deployment reads large files from the volume at +about 104 MiB/s. A model import writes to the same volume at less than 1 MiB/s. + +## Design + +Keep `xio.Copy` as the context-aware copy entry point. Add variadic functional +options so existing callers continue to compile without changes. + +Add an exported `Option` type and a `WithBufferSize(size int) Option` function. +`Copy` uses a 1 MiB buffer by default. A caller can override the buffer size +with `WithBufferSize`. + +If a caller supplies a non-positive buffer size, `Copy` uses the 1 MiB default. +This rule prevents invalid configuration from causing an `io.CopyBuffer` +panic. + +`Copy` allocates one buffer for each active call. It passes that buffer to +`io.CopyBuffer`. The context-aware reader continues to check cancellation +before each source read. + +The first change does not use `sync.Pool`. A pool adds shared state and retains +large caller-selected buffers. Measurements do not justify that complexity. + +## Compatibility + +The existing signature gains only a variadic argument: + +```go +func Copy(ctx context.Context, dst io.Writer, src io.Reader, options ...Option) (int64, error) +``` + +All existing calls remain source compatible. Copy results and cancellation +errors do not change. + +The default buffer increases temporary memory use by approximately 992 KiB for +each concurrent copy compared with the current 32 KiB buffer. + +## Tests and measurement + +Add a Ginkgo suite for `pkg/xio`. Tests cover these behaviors: + +- `Copy` copies the complete source. +- The default buffer permits reads larger than 32 KiB. +- `WithBufferSize` changes the maximum requested read size. +- A non-positive override uses the default buffer. +- A canceled context stops the copy and returns the context error. + +Add a benchmark that runs `Copy` with the default buffer and representative +overrides. The benchmark records throughput and allocations. It does not make +timing assertions. + +Run the focused `pkg/xio` suite first. Then run the packages that call +`xio.Copy`: `pkg/downloader` and `pkg/oci`. + +## Deployment validation + +The code change alone does not alter the running development deployment. After +CI publishes a development image and Flux deploys it, import a large model to +the NAS-backed models directory. Compare the progress rate with the previous +0.7-0.8 MiB/s result. + diff --git a/pkg/xio/copy.go b/pkg/xio/copy.go index 93aaee38a..35b46f33a 100644 --- a/pkg/xio/copy.go +++ b/pkg/xio/copy.go @@ -5,17 +5,39 @@ import ( "io" ) +const defaultBufferSize = 1 << 20 + +type options struct { + bufferSize int +} + +type Option func(*options) + +func WithBufferSize(size int) Option { + return func(options *options) { + if size > 0 { + options.bufferSize = size + } + } +} + type readerFunc func(p []byte) (n int, err error) func (rf readerFunc) Read(p []byte) (n int, err error) { return rf(p) } -func Copy(ctx context.Context, dst io.Writer, src io.Reader) (int64, error) { - return io.Copy(dst, readerFunc(func(p []byte) (int, error) { +func Copy(ctx context.Context, dst io.Writer, src io.Reader, opts ...Option) (int64, error) { + copyOptions := options{bufferSize: defaultBufferSize} + for _, option := range opts { + option(©Options) + } + + buffer := make([]byte, copyOptions.bufferSize) + return io.CopyBuffer(dst, readerFunc(func(p []byte) (int, error) { select { case <-ctx.Done(): return 0, ctx.Err() default: return src.Read(p) } - })) + }), buffer) } diff --git a/pkg/xio/copy_suite_test.go b/pkg/xio/copy_suite_test.go new file mode 100644 index 000000000..859228235 --- /dev/null +++ b/pkg/xio/copy_suite_test.go @@ -0,0 +1,13 @@ +package xio_test + +import ( + "testing" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +func TestXIO(t *testing.T) { + RegisterFailHandler(Fail) + RunSpecs(t, "XIO Suite") +} diff --git a/pkg/xio/copy_test.go b/pkg/xio/copy_test.go new file mode 100644 index 000000000..7d479fe49 --- /dev/null +++ b/pkg/xio/copy_test.go @@ -0,0 +1,115 @@ +package xio_test + +import ( + "bytes" + "context" + "io" + "testing" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/pkg/xio" +) + +type recordingReader struct { + reader io.Reader + maxRead int + reads int +} + +type writerFunc func(p []byte) (int, error) + +func (w writerFunc) Write(p []byte) (int, error) { return w(p) } + +var discardWriter = writerFunc(func(p []byte) (int, error) { return len(p), nil }) + +func (r *recordingReader) Read(p []byte) (int, error) { + r.reads++ + if len(p) > r.maxRead { + r.maxRead = len(p) + } + return r.reader.Read(p) +} + +var _ = Describe("Copy", func() { + It("copies the complete source", func() { + contents := bytes.Repeat([]byte("complete copy"), 10_000) + var destination bytes.Buffer + + written, err := xio.Copy(context.Background(), &destination, bytes.NewReader(contents)) + + Expect(err).NotTo(HaveOccurred()) + Expect(written).To(Equal(int64(len(contents)))) + Expect(destination.Bytes()).To(Equal(contents)) + }) + + It("uses a default read buffer larger than 32 KiB", func() { + source := &recordingReader{reader: bytes.NewReader(make([]byte, 2<<20))} + + _, err := xio.Copy(context.Background(), discardWriter, source) + + Expect(err).NotTo(HaveOccurred()) + Expect(source.maxRead).To(Equal(1 << 20)) + }) + + It("uses a custom buffer size", func() { + const bufferSize = 64 << 10 + source := &recordingReader{reader: bytes.NewReader(make([]byte, 2*bufferSize))} + + _, err := xio.Copy(context.Background(), discardWriter, source, xio.WithBufferSize(bufferSize)) + + Expect(err).NotTo(HaveOccurred()) + Expect(source.maxRead).To(Equal(bufferSize)) + }) + + DescribeTable("falls back to the default buffer for invalid sizes", + func(size int) { + source := &recordingReader{reader: bytes.NewReader(make([]byte, 2<<20))} + + _, err := xio.Copy(context.Background(), discardWriter, source, xio.WithBufferSize(size)) + + Expect(err).NotTo(HaveOccurred()) + Expect(source.maxRead).To(Equal(1 << 20)) + }, + Entry("zero", 0), + Entry("negative", -1), + ) + + It("checks cancellation before reading the source", func() { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + source := &recordingReader{reader: bytes.NewReader([]byte("unread"))} + + written, err := xio.Copy(ctx, io.Discard, source) + + Expect(err).To(MatchError(context.Canceled)) + Expect(written).To(BeZero()) + Expect(source.reads).To(BeZero()) + }) +}) + +func BenchmarkCopy(b *testing.B) { + contents := bytes.Repeat([]byte("benchmark payload"), 1<<16) + tests := []struct { + name string + options []xio.Option + }{ + {name: "default", options: []xio.Option{}}, + {name: "32 KiB", options: []xio.Option{xio.WithBufferSize(32 << 10)}}, + {name: "1 MiB", options: []xio.Option{xio.WithBufferSize(1 << 20)}}, + {name: "4 MiB", options: []xio.Option{xio.WithBufferSize(4 << 20)}}, + } + + for _, test := range tests { + b.Run(test.name, func(b *testing.B) { + b.ReportAllocs() + for b.Loop() { + _, err := xio.Copy(context.Background(), io.Discard, bytes.NewReader(contents), test.options...) + if err != nil { + b.Fatal(err) + } + } + }) + } +}