diff --git a/packages/orchestrator/pkg/nfsproxy/chroot/fs.go b/packages/orchestrator/pkg/nfsproxy/chroot/fs.go index 3960659d67..fb9b861d3c 100644 --- a/packages/orchestrator/pkg/nfsproxy/chroot/fs.go +++ b/packages/orchestrator/pkg/nfsproxy/chroot/fs.go @@ -8,10 +8,13 @@ import ( "github.com/go-git/go-billy/v5" "github.com/e2b-dev/infra/packages/orchestrator/pkg/chrooted" + "github.com/e2b-dev/infra/packages/orchestrator/pkg/nfsproxy/mountcache" ) type wrappedFS struct { - chroot *chrooted.Chrooted + chroot *chrooted.Chrooted + sandboxID string + lifecycleID string } func (f *wrappedFS) Create(filename string) (billy.File, error) { @@ -82,8 +85,12 @@ func (f *wrappedFS) Root() string { return f.chroot.Root() } +func (f *wrappedFS) NFSCacheOwner() mountcache.Owner { + return mountcache.Owner{SandboxID: f.sandboxID, LifecycleID: f.lifecycleID} +} + var _ billy.Filesystem = (*wrappedFS)(nil) -func wrapChrooted(chroot *chrooted.Chrooted) *wrappedFS { - return &wrappedFS{chroot: chroot} +func wrapChrooted(chroot *chrooted.Chrooted, sandboxID, lifecycleID string) *wrappedFS { + return &wrappedFS{chroot: chroot, sandboxID: sandboxID, lifecycleID: lifecycleID} } diff --git a/packages/orchestrator/pkg/nfsproxy/chroot/nfs.go b/packages/orchestrator/pkg/nfsproxy/chroot/nfs.go index 89825a3298..d2d5a4bb5e 100644 --- a/packages/orchestrator/pkg/nfsproxy/chroot/nfs.go +++ b/packages/orchestrator/pkg/nfsproxy/chroot/nfs.go @@ -119,7 +119,19 @@ func (h *NFSHandler) Mount( conn net.Conn, request nfs.MountRequest, ) (nfs.MountStatus, billy.Filesystem, []nfs.AuthFlavor) { - fs, err := h.getChroot(ctx, conn.RemoteAddr(), request) + sbx, err := h.sandboxes.GetByHostPort(conn.RemoteAddr().String()) + if err != nil { + sourceIP, _, _ := net.SplitHostPort(conn.RemoteAddr().String()) + logger.L().Warn(ctx, "failed to get path", + zap.String("request", string(request.Dirpath)), + logger.WithSandboxIP(sourceIP), + zap.Error(fmt.Errorf("%w: %w", ErrUnknownSandbox, err)), + ) + + return nfs.MountStatusErrAcces, mountFailedFS{}, nil + } + + fs, err := h.getChroot(ctx, sbx, request) if err != nil { sourceIP, _, _ := net.SplitHostPort(conn.RemoteAddr().String()) @@ -131,17 +143,12 @@ func (h *NFSHandler) Mount( return nfs.MountStatusErrAcces, mountFailedFS{}, nil } - return nfs.MountStatusOk, wrapChrooted(fs), nil + return nfs.MountStatusOk, wrapChrooted(fs, sbx.Runtime.SandboxID, sbx.LifecycleID), nil } var mountPath = regexp.MustCompile(`^/[^/]+$`) -func (h *NFSHandler) getChroot(ctx context.Context, remoteAddr net.Addr, request nfs.MountRequest) (*chrooted.Chrooted, error) { - sbx, err := h.sandboxes.GetByHostPort(remoteAddr.String()) - if err != nil { - return nil, fmt.Errorf("%w: %w", ErrUnknownSandbox, err) - } - +func (h *NFSHandler) getChroot(ctx context.Context, sbx *sandbox.Sandbox, request nfs.MountRequest) (*chrooted.Chrooted, error) { // normalize the mount path requestedPath := string(request.Dirpath) regexpMatch := mountPath.MatchString(requestedPath) diff --git a/packages/orchestrator/pkg/nfsproxy/mountcache/fs.go b/packages/orchestrator/pkg/nfsproxy/mountcache/fs.go new file mode 100644 index 0000000000..a1795a0e3a --- /dev/null +++ b/packages/orchestrator/pkg/nfsproxy/mountcache/fs.go @@ -0,0 +1,26 @@ +package mountcache + +import ( + "github.com/go-git/go-billy/v5" + "github.com/google/uuid" +) + +type mountedFS struct { + billy.Filesystem + mountID uuid.UUID + owner Owner +} + +func (f *mountedFS) Unwrap() billy.Filesystem { + return f.Filesystem +} + +func (f *mountedFS) NFSCacheMountID() uuid.UUID { + return f.mountID +} + +func (f *mountedFS) NFSCacheOwner() Owner { + return f.owner +} + +var _ billy.Filesystem = (*mountedFS)(nil) diff --git a/packages/orchestrator/pkg/nfsproxy/mountcache/handler.go b/packages/orchestrator/pkg/nfsproxy/mountcache/handler.go new file mode 100644 index 0000000000..72ad888a6d --- /dev/null +++ b/packages/orchestrator/pkg/nfsproxy/mountcache/handler.go @@ -0,0 +1,244 @@ +package mountcache + +import ( + "context" + "net" + "sync" + + "github.com/go-git/go-billy/v5" + "github.com/google/uuid" + "github.com/willscott/go-nfs" + "github.com/willscott/go-nfs/helpers" + "go.opentelemetry.io/otel" + "go.opentelemetry.io/otel/metric" +) + +const uuidLength = len(uuid.UUID{}) + +type mountIdentity interface { + NFSCacheMountID() uuid.UUID +} + +type mountOwner interface { + NFSCacheOwner() Owner +} + +type filesystemUnwrapper interface { + Unwrap() billy.Filesystem +} + +type shard struct { + handler nfs.Handler +} + +type Owner struct { + SandboxID string + LifecycleID string +} + +// Handler gives every successful NFS mount an independent handle cache. Each +// opaque handle is prefixed with its mount ID so it can be routed to the right +// cache without sharing eviction capacity with other mounts. +// +// Handler intentionally does not implement nfs.CachingHandler. Directory +// verifier caching has no filesystem or mount argument and cannot be safely +// isolated by mount. +type Handler struct { + inner nfs.Handler + cacheLimit int + + mu sync.RWMutex + shards map[uuid.UUID]*shard + shardsByOwner map[Owner]map[uuid.UUID]struct{} + shardsGauge metric.Int64ObservableGauge + ownersGauge metric.Int64ObservableGauge +} + +var _ nfs.Handler = (*Handler)(nil) + +func NewHandler(inner nfs.Handler, cacheLimit int) *Handler { + h := &Handler{ + inner: inner, + cacheLimit: cacheLimit, + shards: make(map[uuid.UUID]*shard), + shardsByOwner: make(map[Owner]map[uuid.UUID]struct{}), + } + + meter := otel.Meter("github.com/e2b-dev/infra/packages/orchestrator/pkg/nfsproxy/mountcache") + if gauge, err := meter.Int64ObservableGauge("nfs.mount_cache.shards", metric.WithInt64Callback(func(_ context.Context, observer metric.Int64Observer) error { + h.mu.RLock() + defer h.mu.RUnlock() + + observer.Observe(int64(len(h.shards))) + + return nil + })); err == nil { + h.shardsGauge = gauge + } else { + nfs.Log.Warnf("failed to create NFS mount cache shard gauge: %v", err) + } + + if gauge, err := meter.Int64ObservableGauge("nfs.mount_cache.owners", metric.WithInt64Callback(func(_ context.Context, observer metric.Int64Observer) error { + h.mu.RLock() + defer h.mu.RUnlock() + + observer.Observe(int64(len(h.shardsByOwner))) + + return nil + })); err == nil { + h.ownersGauge = gauge + } else { + nfs.Log.Warnf("failed to create NFS mount cache owner gauge: %v", err) + } + + return h +} + +func (h *Handler) Mount(ctx context.Context, conn net.Conn, request nfs.MountRequest) (nfs.MountStatus, billy.Filesystem, []nfs.AuthFlavor) { + status, filesystem, auth := h.inner.Mount(ctx, conn, request) + if status != nfs.MountStatusOk || filesystem == nil { + return status, filesystem, auth + } + + mountID := uuid.New() + owner, _ := filesystemOwner(filesystem) + + h.mu.Lock() + h.shards[mountID] = &shard{handler: helpers.NewCachingHandler(h.inner, h.cacheLimit)} + if owner.SandboxID != "" && owner.LifecycleID != "" { + if h.shardsByOwner[owner] == nil { + h.shardsByOwner[owner] = make(map[uuid.UUID]struct{}) + } + h.shardsByOwner[owner][mountID] = struct{}{} + } + h.mu.Unlock() + + return status, &mountedFS{Filesystem: filesystem, mountID: mountID, owner: owner}, auth +} + +func (h *Handler) Change(ctx context.Context, filesystem billy.Filesystem) billy.Change { + return h.inner.Change(ctx, filesystem) +} + +func (h *Handler) FSStat(ctx context.Context, filesystem billy.Filesystem, stat *nfs.FSStat) error { + return h.inner.FSStat(ctx, filesystem, stat) +} + +func (h *Handler) ToHandle(ctx context.Context, filesystem billy.Filesystem, path []string) []byte { + mountID, ok := filesystemMountID(filesystem) + if !ok { + return nil + } + + h.mu.RLock() + shard := h.shards[mountID] + h.mu.RUnlock() + if shard == nil { + return nil + } + + localHandle := shard.handler.ToHandle(ctx, filesystem, path) + if len(localHandle) == 0 { + return nil + } + + handle := make([]byte, 0, uuidLength+len(localHandle)) + handle = append(handle, mountID[:]...) + handle = append(handle, localHandle...) + + return handle +} + +func (h *Handler) FromHandle(ctx context.Context, handle []byte) (billy.Filesystem, []string, error) { + shard, localHandle, err := h.resolve(handle) + if err != nil { + return nil, []string{}, err + } + + return shard.handler.FromHandle(ctx, localHandle) +} + +func (h *Handler) InvalidateHandle(ctx context.Context, filesystem billy.Filesystem, handle []byte) error { + shard, localHandle, err := h.resolve(handle) + if err != nil { + return err + } + + return shard.handler.InvalidateHandle(ctx, filesystem, localHandle) +} + +func (h *Handler) HandleLimit() int { + return h.cacheLimit +} + +// RemoveOwner drops all mount caches belonging to one sandbox lifecycle. It is +// safe to call more than once. +func (h *Handler) RemoveOwner(owner Owner) { + if owner.SandboxID == "" || owner.LifecycleID == "" { + return + } + + h.mu.Lock() + defer h.mu.Unlock() + + for mountID := range h.shardsByOwner[owner] { + delete(h.shards, mountID) + } + delete(h.shardsByOwner, owner) +} + +func (h *Handler) resolve(handle []byte) (*shard, []byte, error) { + if len(handle) <= uuidLength { + return nil, nil, staleHandle() + } + + mountID, err := uuid.FromBytes(handle[:uuidLength]) + if err != nil { + return nil, nil, staleHandle() + } + + h.mu.RLock() + shard := h.shards[mountID] + h.mu.RUnlock() + if shard == nil { + return nil, nil, staleHandle() + } + + return shard, handle[uuidLength:], nil +} + +func staleHandle() error { + return &nfs.NFSStatusError{NFSStatus: nfs.NFSStatusStale} +} + +func filesystemMountID(filesystem billy.Filesystem) (uuid.UUID, bool) { + for filesystem != nil { + if identity, ok := filesystem.(mountIdentity); ok { + return identity.NFSCacheMountID(), true + } + + unwrapper, ok := filesystem.(filesystemUnwrapper) + if !ok { + return uuid.Nil, false + } + filesystem = unwrapper.Unwrap() + } + + return uuid.Nil, false +} + +func filesystemOwner(filesystem billy.Filesystem) (Owner, bool) { + for filesystem != nil { + if owner, ok := filesystem.(mountOwner); ok { + return owner.NFSCacheOwner(), true + } + + unwrapper, ok := filesystem.(filesystemUnwrapper) + if !ok { + return Owner{}, false + } + filesystem = unwrapper.Unwrap() + } + + return Owner{}, false +} diff --git a/packages/orchestrator/pkg/nfsproxy/mountcache/handler_test.go b/packages/orchestrator/pkg/nfsproxy/mountcache/handler_test.go new file mode 100644 index 0000000000..5027d702d4 --- /dev/null +++ b/packages/orchestrator/pkg/nfsproxy/mountcache/handler_test.go @@ -0,0 +1,128 @@ +package mountcache + +import ( + "context" + "net" + "testing" + + "github.com/go-git/go-billy/v5" + "github.com/go-git/go-billy/v5/memfs" + "github.com/stretchr/testify/require" + "github.com/willscott/go-nfs" +) + +type testFilesystem struct { + billy.Filesystem + owner Owner +} + +func (f *testFilesystem) NFSCacheOwner() Owner { + return f.owner +} + +type testHandler struct { + lifecycles []string +} + +func (h *testHandler) Mount(context.Context, net.Conn, nfs.MountRequest) (nfs.MountStatus, billy.Filesystem, []nfs.AuthFlavor) { + lifecycleID := "lifecycle-a" + if len(h.lifecycles) > 0 { + lifecycleID = h.lifecycles[0] + h.lifecycles = h.lifecycles[1:] + } + + filesystem := &testFilesystem{ + Filesystem: memfs.New(), + owner: Owner{SandboxID: "sandbox-a", LifecycleID: lifecycleID}, + } + + return nfs.MountStatusOk, filesystem, nil +} + +func (*testHandler) Change(context.Context, billy.Filesystem) billy.Change { return nil } +func (*testHandler) FSStat(context.Context, billy.Filesystem, *nfs.FSStat) error { return nil } +func (*testHandler) ToHandle(context.Context, billy.Filesystem, []string) []byte { + panic("not reached") +} +func (*testHandler) FromHandle(context.Context, []byte) (billy.Filesystem, []string, error) { + panic("not reached") +} +func (*testHandler) InvalidateHandle(context.Context, billy.Filesystem, []byte) error { + panic("not reached") +} +func (*testHandler) HandleLimit() int { return 1024 } + +func requireStaleHandle(t *testing.T, err error) { + t.Helper() + + var statusErr *nfs.NFSStatusError + require.ErrorAs(t, err, &statusErr) + require.Equal(t, nfs.NFSStatusStale, statusErr.NFSStatus) +} + +func TestMountCreatesIndependentCaches(t *testing.T) { + t.Parallel() + + handler := NewHandler(&testHandler{}, 2) + _, filesystemA, _ := handler.Mount(t.Context(), nil, nfs.MountRequest{}) + _, filesystemB, _ := handler.Mount(t.Context(), nil, nfs.MountRequest{}) + + handleA := handler.ToHandle(t.Context(), filesystemA, nil) + handleB := handler.ToHandle(t.Context(), filesystemB, nil) + require.Len(t, handleA, 2*uuidLength) + require.Len(t, handleB, 2*uuidLength) + require.NotEqual(t, handleA[:uuidLength], handleB[:uuidLength]) + + handler.ToHandle(t.Context(), filesystemA, []string{"one"}) + handler.ToHandle(t.Context(), filesystemA, []string{"two"}) + _, _, err := handler.FromHandle(t.Context(), handleA) + requireStaleHandle(t, err) + + resolved, path, err := handler.FromHandle(t.Context(), handleB) + require.NoError(t, err) + require.Same(t, filesystemB, resolved) + require.Empty(t, path) +} + +func TestRemoveOwnerInvalidatesAllMounts(t *testing.T) { + t.Parallel() + + handler := NewHandler(&testHandler{}, 1024) + _, filesystemA, _ := handler.Mount(t.Context(), nil, nfs.MountRequest{}) + _, filesystemB, _ := handler.Mount(t.Context(), nil, nfs.MountRequest{}) + handleA := handler.ToHandle(t.Context(), filesystemA, nil) + handleB := handler.ToHandle(t.Context(), filesystemB, nil) + + handler.RemoveOwner(Owner{SandboxID: "sandbox-a", LifecycleID: "lifecycle-a"}) + + _, _, err := handler.FromHandle(t.Context(), handleA) + requireStaleHandle(t, err) + _, _, err = handler.FromHandle(t.Context(), handleB) + requireStaleHandle(t, err) +} + +func TestRemoveOwnerOnlyInvalidatesMatchingLifecycle(t *testing.T) { + t.Parallel() + + handler := NewHandler(&testHandler{lifecycles: []string{"lifecycle-old", "lifecycle-new"}}, 1024) + _, oldFilesystem, _ := handler.Mount(t.Context(), nil, nfs.MountRequest{}) + _, newFilesystem, _ := handler.Mount(t.Context(), nil, nfs.MountRequest{}) + oldHandle := handler.ToHandle(t.Context(), oldFilesystem, nil) + newHandle := handler.ToHandle(t.Context(), newFilesystem, nil) + + handler.RemoveOwner(Owner{SandboxID: "sandbox-a", LifecycleID: "lifecycle-old"}) + + _, _, err := handler.FromHandle(t.Context(), oldHandle) + requireStaleHandle(t, err) + resolved, _, err := handler.FromHandle(t.Context(), newHandle) + require.NoError(t, err) + require.Same(t, newFilesystem, resolved) +} + +func TestHandlerDoesNotExposeVerifierCaching(t *testing.T) { + t.Parallel() + + handler := NewHandler(&testHandler{}, 1024) + _, implementsCachingHandler := any(handler).(nfs.CachingHandler) + require.False(t, implementsCachingHandler) +} diff --git a/packages/orchestrator/pkg/nfsproxy/proxy.go b/packages/orchestrator/pkg/nfsproxy/proxy.go index cb85da12a8..34ac1443bf 100644 --- a/packages/orchestrator/pkg/nfsproxy/proxy.go +++ b/packages/orchestrator/pkg/nfsproxy/proxy.go @@ -10,13 +10,13 @@ import ( "sync" "github.com/willscott/go-nfs" - "github.com/willscott/go-nfs/helpers" "github.com/e2b-dev/infra/packages/orchestrator/pkg/chrooted" "github.com/e2b-dev/infra/packages/orchestrator/pkg/nfsproxy/cfg" "github.com/e2b-dev/infra/packages/orchestrator/pkg/nfsproxy/chroot" "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/mountcache" "github.com/e2b-dev/infra/packages/orchestrator/pkg/nfsproxy/recovery" "github.com/e2b-dev/infra/packages/orchestrator/pkg/nfsproxy/tracing" "github.com/e2b-dev/infra/packages/orchestrator/pkg/sandbox" @@ -31,6 +31,19 @@ type Proxy struct { server *nfs.Server } +type cacheLifecycleSubscriber struct { + handler *mountcache.Handler +} + +func (s cacheLifecycleSubscriber) OnInsert(_ context.Context, _ *sandbox.Sandbox) {} + +func (s cacheLifecycleSubscriber) OnNetworkRelease(_ context.Context, sbx *sandbox.Sandbox) { + s.handler.RemoveOwner(mountcache.Owner{ + SandboxID: sbx.Runtime.SandboxID, + LifecycleID: sbx.LifecycleID, + }) +} + func NewProxy(ctx context.Context, builder *chrooted.Builder, sandboxes *sandbox.Map, config cfg.Config) (*Proxy, error) { setLogLevelOnce.Do(func() { nfs.Log.SetLevel(config.NFSLogLevel) @@ -47,7 +60,9 @@ func NewProxy(ctx context.Context, builder *chrooted.Builder, sandboxes *sandbox } // wrap the handler in middleware - handler = helpers.NewCachingHandler(handler, cacheLimit) + cachingHandler := mountcache.NewHandler(handler, cacheLimit) + sandboxes.Subscribe(cacheLifecycleSubscriber{handler: cachingHandler}) + handler = cachingHandler if config.Tracing { handler = tracing.WrapWithTracing(handler, config)