(buildkit proxy) fix: double memory usage during commit (#54)

* refactor cache backend client and add logs

* run workflow to test

* disable proxy

* systemd stuff

* path fixes

* halve memory usage

* push to run

* url fixes

* rm test stuff from this repo

* release dev

* compile time check

* add logs

* too many logs

* better logs

* add tests

* more logs

* fix switch cases

* simplify chunk reader logic

* fix error handling

* rm logs

Author: abhijit-hota

.github/workflows/release-testing.yaml

@@ -3,7 +3,7 @@ name: Branch Build and Upload to R2
 on:
   push:
     branches:
-      - "prashant-warp-740"
+      - ah-docker-layer-caching
 
 env:
   AWS_ACCESS_KEY_ID: ${{ secrets.WB_PACKAGES_DEV_R2_ACCESS_KEY_ID }}
@@ -26,7 +26,7 @@ jobs:
       - name: Set up Go
         uses: WarpBuilds/setup-go@v5
         with:
-          go-version: "1.25"
+          go-version: "1.24"
           cache: true
 
       - name: Define S3 Path

go.mod

@@ -1,8 +1,6 @@
 module github.com/warpbuilds/warpbuild-agent
 
-go 1.23.0
-
-toolchain go1.24.1
+go 1.24
 
 require (
 	cloud.google.com/go/storage v1.43.0

pkg/app/app.go

@@ -238,9 +238,11 @@ func NewApp(ctx context.Context, opts *ApplicationOptions) error {
 
 	} else if opts.LaunchProxyServer {
 		proxy.StartProxyServer(ctx, &proxy.ProxyServerOptions{
-			CacheBackendHost:                 settings.Proxy.CacheBackendHost,
-			CacheProxyPort:                   settings.Proxy.CacheProxyPort,
-			WarpBuildRunnerVerificationToken: settings.Agent.RunnerVerificationToken,
+			Port: settings.Proxy.CacheProxyPort,
+			CacheBackendInfo: proxy.CacheBackendInfo{
+				HostURL:   settings.Proxy.CacheBackendHost,
+				AuthToken: settings.Agent.RunnerVerificationToken,
+			},
 		})
 	} else if opts.LaunchTransparentCache {
 		// Start the transparent cache server with configured ports

pkg/proxy/cache_backend_client.go

@@ -0,0 +1,57 @@
+package proxy
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"net/url"
+
+	"github.com/gofiber/fiber/v2"
+)
+
+type CacheBackendRequest struct {
+	Path string
+	Body any
+}
+
+func getCacheBackendInfo(ctx context.Context) CacheBackendInfo {
+	opts := ctx.Value(PROXY_SERVER_OPTIONS_CONTEXT_KEY).(*ProxyServerOptions)
+	return opts.CacheBackendInfo
+}
+
+func callCacheBackend[T any](ctx context.Context, req CacheBackendRequest) (*T, error) {
+	info := getCacheBackendInfo(ctx)
+	requestURL, err := url.JoinPath(info.HostURL, "/v1/cache", req.Path)
+	if err != nil {
+		return nil, fmt.Errorf("failed to construct URL: %w", err)
+	}
+
+	f := fiber.Post(requestURL).
+		ContentType("application/json").
+		Add("Accept", "application/json").
+		Add("Authorization", fmt.Sprintf("Bearer %s", info.AuthToken))
+
+	if req.Body != nil {
+		bodyBytes, err := json.Marshal(req.Body)
+		if err != nil {
+			return nil, fmt.Errorf("failed to marshal request body: %w", err)
+		}
+		f = f.Body(bodyBytes)
+	}
+
+	statusCode, body, errs := f.Bytes()
+	if len(errs) > 0 {
+		return nil, fmt.Errorf("request failed: %v", errs)
+	}
+
+	if statusCode < 200 || statusCode >= 300 {
+		return nil, fmt.Errorf("request failed with status %d: %s", statusCode, string(body))
+	}
+
+	var result T
+	if err := json.Unmarshal(body, &result); err != nil {
+		return nil, fmt.Errorf("failed to parse response: %w", err)
+	}
+
+	return &result, nil
+}

pkg/proxy/chunk_reader.go

@@ -0,0 +1,30 @@
+package proxy
+
+import (
+	"bytes"
+	"io"
+	"sort"
+)
+
+func NewUnorderedChunkReader(chunksMap map[int64]ChunkData) (io.Reader, int64) {
+	// Sort the chunks by start offset
+	offsets := make([]int64, 0, len(chunksMap))
+	for offset := range chunksMap {
+		offsets = append(offsets, offset)
+	}
+	sort.Slice(offsets, func(i, j int) bool {
+		return chunksMap[offsets[i]].StartOffset < chunksMap[offsets[j]].StartOffset
+	})
+
+	// Create readers for each chunk in order
+	readers := make([]io.Reader, len(offsets))
+	var totalSize int64
+	for i, offset := range offsets {
+		chunk := chunksMap[offset]
+		readers[i] = bytes.NewReader(chunk.Content)
+		totalSize += int64(len(chunk.Content))
+	}
+
+	// Chain them together with stdlib MultiReader
+	return io.MultiReader(readers...), totalSize
+}