nfsproxy: preserve caching handler through middleware chain

packages/orchestrator/pkg/nfsproxy/logged/handler.go

@@ -3,6 +3,7 @@ package logged
 import (
 	"context"
 	"fmt"
+	iofs "io/fs"
 	"net"
 
 	"github.com/go-git/go-billy/v5"
@@ -17,7 +18,10 @@ type loggedHandler struct {
 	config cfg.Config
 }
 
-var _ nfs.Handler = (*loggedHandler)(nil)
+var (
+	_ nfs.Handler        = (*loggedHandler)(nil)
+	_ nfs.CachingHandler = (*loggedHandler)(nil)
+)
 
 func WrapWithLogging(ctx context.Context, handler nfs.Handler, config cfg.Config) nfs.Handler {
 	return loggedHandler{ctx: ctx, inner: handler, config: config}
@@ -96,3 +100,21 @@ func (e loggedHandler) HandleLimit() int {
 
 	return e.inner.HandleLimit()
 }
+
+func (e loggedHandler) VerifierFor(path string, contents []iofs.FileInfo) uint64 {
+	handler, ok := e.inner.(nfs.CachingHandler)
+	if !ok {
+		return 0
+	}
+
+	return handler.VerifierFor(path, contents)
+}
+
+func (e loggedHandler) DataForVerifier(path string, verifier uint64) []iofs.FileInfo {
+	handler, ok := e.inner.(nfs.CachingHandler)
+	if !ok {
+		return nil
+	}
+
+	return handler.DataForVerifier(path, verifier)
+}

packages/orchestrator/pkg/nfsproxy/metrics/handler.go

@@ -3,6 +3,7 @@ package metrics
 import (
 	"context"
 	"fmt"
+	iofs "io/fs"
 	"net"
 
 	"github.com/go-git/go-billy/v5"
@@ -16,7 +17,10 @@ type metricsHandler struct {
 	config cfg.Config
 }
 
-var _ nfs.Handler = (*metricsHandler)(nil)
+var (
+	_ nfs.Handler        = (*metricsHandler)(nil)
+	_ nfs.CachingHandler = (*metricsHandler)(nil)
+)
 
 func WrapWithMetrics(handler nfs.Handler, config cfg.Config) nfs.Handler {
 	return &metricsHandler{inner: handler, config: config}
@@ -98,3 +102,21 @@ func (m *metricsHandler) InvalidateHandle(ctx context.Context, filesystem billy.
 func (m *metricsHandler) HandleLimit() int {
 	return m.inner.HandleLimit()
 }
+
+func (m *metricsHandler) VerifierFor(path string, contents []iofs.FileInfo) uint64 {
+	handler, ok := m.inner.(nfs.CachingHandler)
+	if !ok {
+		return 0
+	}
+
+	return handler.VerifierFor(path, contents)
+}
+
+func (m *metricsHandler) DataForVerifier(path string, verifier uint64) []iofs.FileInfo {
+	handler, ok := m.inner.(nfs.CachingHandler)
+	if !ok {
+		return nil
+	}
+
+	return handler.DataForVerifier(path, verifier)
+}

packages/orchestrator/pkg/nfsproxy/recovery/cache_passthrough_test.go

@@ -0,0 +1,62 @@
+package recovery
+
+import (
+	iofs "io/fs"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	nfs "github.com/willscott/go-nfs"
+
+	"github.com/e2b-dev/infra/packages/orchestrator/pkg/nfsproxy/cfg"
+	"github.com/e2b-dev/infra/packages/orchestrator/pkg/nfsproxy/logged"
+	"github.com/e2b-dev/infra/packages/orchestrator/pkg/nfsproxy/metrics"
+	"github.com/e2b-dev/infra/packages/orchestrator/pkg/nfsproxy/tracing"
+)
+
+type cachingHandlerStub struct {
+	nfs.Handler
+
+	verifier         uint64
+	contents         []iofs.FileInfo
+	verifierPath     string
+	verifierContents []iofs.FileInfo
+	dataPath         string
+	dataVerifier     uint64
+}
+
+func (h *cachingHandlerStub) VerifierFor(path string, contents []iofs.FileInfo) uint64 {
+	h.verifierPath = path
+	h.verifierContents = contents
+
+	return h.verifier
+}
+
+func (h *cachingHandlerStub) DataForVerifier(path string, verifier uint64) []iofs.FileInfo {
+	h.dataPath = path
+	h.dataVerifier = verifier
+
+	return h.contents
+}
+
+func TestCachingHandlerPassthroughThroughMiddlewareChain(t *testing.T) {
+	t.Parallel()
+
+	contents := []iofs.FileInfo{nil}
+	inner := &cachingHandlerStub{verifier: 42, contents: contents}
+	config := cfg.Config{}
+
+	var handler nfs.Handler = inner
+	handler = tracing.WrapWithTracing(handler, config)
+	handler = metrics.WrapWithMetrics(handler, config)
+	handler = logged.WrapWithLogging(t.Context(), handler, config)
+	handler = WrapWithRecovery(t.Context(), handler)
+
+	cachingHandler, ok := handler.(nfs.CachingHandler)
+	require.True(t, ok)
+	require.Equal(t, uint64(42), cachingHandler.VerifierFor("/dir", contents))
+	require.Equal(t, contents, cachingHandler.DataForVerifier("/dir", 42))
+	require.Equal(t, "/dir", inner.verifierPath)
+	require.Equal(t, contents, inner.verifierContents)
+	require.Equal(t, "/dir", inner.dataPath)
+	require.Equal(t, uint64(42), inner.dataVerifier)
+}

packages/orchestrator/pkg/nfsproxy/recovery/main.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	iofs "io/fs"
 	"net"
 
 	"github.com/go-git/go-billy/v5"
@@ -18,7 +19,10 @@ type Handler struct {
 	ctx   context.Context //nolint:containedctx // can't change the API, still need it
 }
 
-var _ nfs.Handler = (*Handler)(nil)
+var (
+	_ nfs.Handler        = (*Handler)(nil)
+	_ nfs.CachingHandler = (*Handler)(nil)
+)
 
 func WrapWithRecovery(ctx context.Context, h nfs.Handler) *Handler {
 	return &Handler{inner: h, ctx: ctx}
@@ -69,6 +73,28 @@ func (h Handler) HandleLimit() int {
 	return h.inner.HandleLimit()
 }
 
+func (h Handler) VerifierFor(path string, contents []iofs.FileInfo) (verifier uint64) {
+	defer tryRecovery(h.ctx, "VerifierFor")
+
+	handler, ok := h.inner.(nfs.CachingHandler)
+	if !ok {
+		return 0
+	}
+
+	return handler.VerifierFor(path, contents)
+}
+
+func (h Handler) DataForVerifier(path string, verifier uint64) (contents []iofs.FileInfo) {
+	defer tryRecovery(h.ctx, "DataForVerifier")
+
+	handler, ok := h.inner.(nfs.CachingHandler)
+	if !ok {
+		return nil
+	}
+
+	return handler.DataForVerifier(path, verifier)
+}
+
 // tryRecovery must be called via `defer` directly so that recover() runs in
 // the deferred function frame. Nesting it inside another helper would cause
 // recover() to return nil and the panic would propagate.

packages/orchestrator/pkg/nfsproxy/tracing/handler.go

@@ -3,6 +3,7 @@ package tracing
 import (
 	"context"
 	"fmt"
+	iofs "io/fs"
 	"net"
 
 	"github.com/go-git/go-billy/v5"
@@ -17,7 +18,10 @@ type tracingHandler struct {
 	config cfg.Config
 }
 
-var _ nfs.Handler = (*tracingHandler)(nil)
+var (
+	_ nfs.Handler        = (*tracingHandler)(nil)
+	_ nfs.CachingHandler = (*tracingHandler)(nil)
+)
 
 func WrapWithTracing(handler nfs.Handler, config cfg.Config) nfs.Handler {
 	return &tracingHandler{inner: handler, config: config}
@@ -101,3 +105,21 @@ func (e *tracingHandler) InvalidateHandle(ctx context.Context, filesystem billy.
 func (e *tracingHandler) HandleLimit() int {
 	return e.inner.HandleLimit()
 }
+
+func (e *tracingHandler) VerifierFor(path string, contents []iofs.FileInfo) uint64 {
+	handler, ok := e.inner.(nfs.CachingHandler)
+	if !ok {
+		return 0
+	}
+
+	return handler.VerifierFor(path, contents)
+}
+
+func (e *tracingHandler) DataForVerifier(path string, verifier uint64) []iofs.FileInfo {
+	handler, ok := e.inner.(nfs.CachingHandler)
+	if !ok {
+		return nil
+	}
+
+	return handler.DataForVerifier(path, verifier)
+}