From 8fd78011a8a42ec8ebacbea0d37f31b6d4e784f1 Mon Sep 17 00:00:00 2001 From: Stephan Renatus Date: Fri, 17 Mar 2023 10:48:36 +0100 Subject: [PATCH] download: copy bundle into buffer only if needed (#5767) The need would be if `d.persist` is true: we're writing it into a file in a later code path. However, if we don't have bundle persistency enabled, this should allow us to get rid of one copy of the bundle in memory. To achieve that, the mechanism used to figure out the bundle size is replaced: - Before, we'd write it into a buffer that we might not need, and use its Len() - After, we're writing it into a no-op io.Writer that only keeps track of the bytes read. Signed-off-by: Stephan Renatus --- download/download.go | 27 +++++++++++++++++++++++++-- 1 file changed, 25 insertions(+), 2 deletions(-) diff --git a/download/download.go b/download/download.go index bc03f8448e..9c611bb647 100644 --- a/download/download.go +++ b/download/download.go @@ -315,7 +315,16 @@ func (d *Downloader) download(ctx context.Context, m metrics.Metrics) (*download defer m.Timer(metrics.RegoLoadBundles).Stop() baseURL := path.Join(d.client.Config().URL, d.path) - loader := bundle.NewTarballLoaderWithBaseURL(io.TeeReader(resp.Body, &buf), baseURL) + cnt := &count{} + r := io.TeeReader(resp.Body, cnt) + + var loader bundle.DirectoryLoader + if d.persist { + tee := io.TeeReader(r, &buf) + loader = bundle.NewTarballLoaderWithBaseURL(tee, baseURL) + } else { + loader = bundle.NewTarballLoaderWithBaseURL(r, baseURL) + } etag := resp.Header.Get("ETag") @@ -356,7 +365,7 @@ func (d *Downloader) download(ctx context.Context, m metrics.Metrics) (*download raw: &buf, etag: etag, longPoll: isLongPollSupported(resp.Header), - size: buf.Len(), + size: cnt.Bytes(), }, nil } @@ -383,6 +392,20 @@ func (d *Downloader) download(ctx context.Context, m metrics.Metrics) (*download } } +type count struct { + total int +} + +func (c *count) Write(p []byte) (n int, err error) { + n = len(p) + c.total += n + return +} + +func (c *count) Bytes() int { + return c.total +} + func isLongPollSupported(header http.Header) bool { return header.Get("Content-Type") == "application/vnd.openpolicyagent.bundles" }