diff --git a/.github/workflows/pr-tests.yml b/.github/workflows/pr-tests.yml index d3f3973974..330b963bf8 100644 --- a/.github/workflows/pr-tests.yml +++ b/.github/workflows/pr-tests.yml @@ -46,9 +46,19 @@ jobs: sudo mount -t hugetlbfs none /mnt/hugepages echo 128 | sudo tee /proc/sys/vm/nr_hugepages + - name: Enable NBD + run: | + sudo modprobe nbd nbds_max=256 + + - name: Disable inotify watching of change events for NBD devices + run: | + echo 'ACTION=="add|change", KERNEL=="nbd*", OPTIONS:="nowatch"' | sudo tee /etc/udev/rules.d/97-nbd-device.rules + sudo udevadm control --reload-rules + sudo udevadm trigger + - name: Run tests working-directory: ${{ matrix.package }} - run: go test -v ${{ matrix.test_path }} + run: sudo go test -v ${{ matrix.test_path }} validate-iac: name: Validate terraform diff --git a/packages/orchestrator/Makefile b/packages/orchestrator/Makefile index 8d1e8d0bfb..b620e92026 100644 --- a/packages/orchestrator/Makefile +++ b/packages/orchestrator/Makefile @@ -85,10 +85,6 @@ build-and-upload/template-manager: build upload/template-manager mock: sudo TEMPLATE_BUCKET_NAME=$(TEMPLATE_BUCKET_NAME) CONSUL_TOKEN=$(CONSUL_TOKEN) NODE_ID="testclient" go run cmd/mock-sandbox/mock.go -template 5wzg6c91u51yaebviysf -build "f0370054-b669-eeee-b33b-573d5287c6ef" -alive 1 -count 2 -.PHONY: mock-nbd -mock-nbd: - sudo go run -gcflags=all="-N -l" cmd/mock-nbd/mock.go - .PHONY: mock-snapshot mock-snapshot: sudo TEMPLATE_BUCKET_NAME=$(TEMPLATE_BUCKET_NAME) CONSUL_TOKEN=$(CONSUL_TOKEN) NODE_ID="testclient" go run cmd/mock-snapshot/mock.go -template 5wzg6c91u51yaebviysf -build "f0370054-b669-4d7e-b33b-573d5287c6ef" -alive 1 -count 1 diff --git a/packages/orchestrator/cmd/mock-nbd/mock.go b/packages/orchestrator/cmd/mock-nbd/mock.go deleted file mode 100644 index e1f300d4a9..0000000000 --- a/packages/orchestrator/cmd/mock-nbd/mock.go +++ /dev/null @@ -1,202 +0,0 @@ -package main - -import ( - "bytes" - "context" - "crypto/rand" - "fmt" - "os" - "os/signal" - - "github.com/google/uuid" - "github.com/pojntfx/go-nbd/pkg/backend" - - "github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/block" - "github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/nbd" - "github.com/e2b-dev/infra/packages/shared/pkg/storage/header" -) - -const blockSize = 4096 - -type DeviceWithClose struct { - b backend.Backend -} - -var _ block.Device = (*DeviceWithClose)(nil) - -func (d *DeviceWithClose) ReadAt(_ context.Context, p []byte, off int64) (n int, err error) { - return d.b.ReadAt(p, off) -} - -func (d *DeviceWithClose) Size() (int64, error) { - return d.b.Size() -} - -func (d *DeviceWithClose) WriteAt(p []byte, off int64) (n int, err error) { - return d.b.WriteAt(p, off) -} - -func (d *DeviceWithClose) Close() error { - return nil -} - -func (d *DeviceWithClose) Slice(_ context.Context, offset, length int64) ([]byte, error) { - b := make([]byte, length) - - _, err := d.b.ReadAt(b, offset) - if err != nil { - return nil, err - } - - return b, nil -} - -func (d *DeviceWithClose) BlockSize() int64 { - return blockSize -} - -func (d *DeviceWithClose) Header() *header.Header { - size, err := d.b.Size() - if err != nil { - panic(err) - } - - h, err := header.NewHeader(header.NewTemplateMetadata( - uuid.New(), - uint64(blockSize), - uint64(size), - ), nil) - if err != nil { - panic(err) - } - - return h -} - -func main() { - data := make([]byte, blockSize*8) - rand.Read(data) - - device := &DeviceWithClose{ - b: backend.NewMemoryBackend(data), - } - - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - done := make(chan os.Signal, 1) - signal.Notify(done, os.Interrupt) - devicePool, err := nbd.NewDevicePool() - if err != nil { - fmt.Fprintf(os.Stderr, "failed to create device pool: %v\n", err) - - return - } - go func() { - devicePool.Populate(ctx) - fmt.Fprintf(os.Stderr, "device pool done populating\n") - }() - defer func() { - err = devicePool.Close(ctx) - if err != nil { - fmt.Fprintf(os.Stderr, "failed to close device pool: %v\n", err) - - return - } - }() - - go func() { - <-done - - cancel() - }() - - for i := 0; ; i++ { - select { - case <-ctx.Done(): - return - default: - } - fmt.Printf("----------------------------------------\n") - fmt.Printf("[%d] starting mock nbd server\n", i) - - readData, err := MockNbd(ctx, device, i, devicePool) - if err != nil { - fmt.Fprintf(os.Stderr, "[%d] failed to mock nbd: %v\n", i, err) - - return - } - - if !bytes.Equal(data, readData) { - fmt.Fprintf(os.Stderr, "[%d] data mismatch\n", i) - - return - } - } -} - -func MockNbd(ctx context.Context, device *DeviceWithClose, index int, devicePool *nbd.DevicePool) ([]byte, error) { - ctx, cancel := context.WithCancel(ctx) - defer cancel() - - size, err := device.Size() - if err != nil { - return nil, fmt.Errorf("failed to get size: %w", err) - } - - deviceIndex, err := devicePool.GetDevice(ctx) - if err != nil { - return nil, fmt.Errorf("failed to get device: %w", err) - } - - var mnt *nbd.DirectPathMount - - defer func() { - counter := 0 - - for { - counter++ - err = devicePool.ReleaseDevice(ctx, deviceIndex) - if err != nil { - if counter%10 == 0 { - fmt.Printf("[%d] failed to release device: %v\n", index, err) - } - - if mnt != nil { - mnt.Close(ctx) - } - - continue - } - - fmt.Printf("[%d] released device: %d\n", index, deviceIndex) - - return - } - }() - - mnt = nbd.NewDirectPathMount(device, devicePool) - - go func() { - <-ctx.Done() - - mnt.Close(context.Background()) //nolint:contextcheck // TODO: fix this later - }() - - _, err = mnt.Open(ctx) - if err != nil { - return nil, fmt.Errorf("failed to open: %w", err) - } - - data := make([]byte, size) - _, err = mnt.Backend.ReadAt(ctx, data, 0) - if err != nil { - return nil, fmt.Errorf("failed to read: %w", err) - } - - fmt.Printf("[%d] Read %d bytes from nbd\n", index, len(data)) - - cancel() - - return data, nil -} diff --git a/packages/orchestrator/cmd/mount-rootfs/main.go b/packages/orchestrator/cmd/mount-rootfs/main.go index e67d95f8a6..b66b90191c 100644 --- a/packages/orchestrator/cmd/mount-rootfs/main.go +++ b/packages/orchestrator/cmd/mount-rootfs/main.go @@ -4,76 +4,49 @@ import ( "context" "flag" "fmt" + "log" "os" + "os/exec" "os/signal" "path/filepath" "time" "github.com/google/uuid" - "go.opentelemetry.io/otel/metric/noop" - "golang.org/x/sys/unix" + "go.uber.org/zap" - "github.com/e2b-dev/infra/packages/orchestrator/internal/cfg" "github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/block" - "github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/block/metrics" - "github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/build" - "github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/nbd" - featureflags "github.com/e2b-dev/infra/packages/shared/pkg/feature-flags" - "github.com/e2b-dev/infra/packages/shared/pkg/storage" - "github.com/e2b-dev/infra/packages/shared/pkg/storage/header" + "github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/nbd/testutils" ) -var _ block.ReadonlyDevice = (*mountReadonlyDevice)(nil) - -type mountReadonlyDevice struct { - *build.File - - header *header.Header - blockSize int64 -} - -func newReadonlyDevice(file *build.File, header *header.Header, blockSize int64) *mountReadonlyDevice { - return &mountReadonlyDevice{ - File: file, - header: header, - blockSize: blockSize, - } -} - -func (m *mountReadonlyDevice) Close() error { - return nil -} - -func (m *mountReadonlyDevice) BlockSize() int64 { - return m.blockSize -} - -func (m *mountReadonlyDevice) Header() *header.Header { - return m.header -} - -func (m *mountReadonlyDevice) Size() (int64, error) { - return int64(m.header.Metadata.Size), nil -} - func main() { - buildId := flag.String("build", "", "build id") - mountPath := flag.String("mount", "", "mount path") + buildId := flag.String("build", "", "build id (only used when empty flag is false)") + mountPath := flag.String("mount", "", "mount path (only used when empty flag is false)") + verify := flag.Bool("verify", false, "verify rootfs integrity (only used when empty flag is false)") + logging := flag.Bool("log", false, "enable logging (it is pretty spammy)") + empty := flag.Bool("empty", false, "create an empty rootfs") + size := flag.Int64("size", 1024*1024*1024, "size of the rootfs (only used when empty flag is true)") + blockSize := flag.Int64("block-size", 4096, "block size of the rootfs (only used when empty flag is true)") flag.Parse() + if *verify && *mountPath == "" { + log.Fatalf("verify flag is only supported when mount path is provided") + } + ctx, cancel := context.WithCancel(context.Background()) defer cancel() done := make(chan os.Signal, 1) signal.Notify(done, os.Interrupt) - // Disabling the logger for normal use—is very spammy, because Populate on device pool periodically logs errors if the number of acquirable devices is less than the number of requested devices. - // logger, err := zap.NewDevelopment() - // if err != nil { - // log.Fatalf("failed to create logger: %s", err) - // } - // zap.ReplaceGlobals(logger) + // Logger is very spammy, because Populate on device pool periodically logs errors if the number of acquirable devices is less than the number of requested devices. + if *logging { + logger, err := zap.NewDevelopment() + if err != nil { + panic(fmt.Errorf("failed to create logger: %w", err)) + } + zap.ReplaceGlobals(logger) + } go func() { <-done @@ -84,103 +57,73 @@ func main() { // We use a separate ctx for majority of the operations as cancelling context for the NBD+storage and *then* doing cleanup for these often resulted in deadlocks. nbdContext := context.Background() - err := mountRootfs(ctx, nbdContext, *buildId, *mountPath) - if err != nil { - panic(fmt.Errorf("failed to mount rootfs: %w", err)) - } -} - -func mountRootfs(ctx, nbdContext context.Context, buildID, mountPath string) error { - files := storage.TemplateFiles{ - BuildID: buildID, - } - - s, err := storage.GetTemplateStorageProvider(nbdContext, nil) - if err != nil { - return fmt.Errorf("failed to get storage provider: %w", err) - } - - obj, err := s.OpenObject(nbdContext, files.StorageRootfsHeaderPath(), storage.RootFSHeaderObjectType) - if err != nil { - return fmt.Errorf("failed to open object: %w", err) - } - - h, err := header.Deserialize(nbdContext, obj) - if err != nil { - id, err := uuid.Parse(buildID) - if err != nil { - return fmt.Errorf("failed to parse build id: %w", err) - } - - r, err := s.OpenSeekableObject(nbdContext, files.StorageRootfsPath(), storage.RootFSObjectType) + if *empty { + err := runEmpty(ctx, nbdContext, *size, *blockSize) if err != nil { - return fmt.Errorf("failed to open object: %w", err) + panic(fmt.Errorf("failed to create empty rootfs: %w", err)) } - - size, err := r.Size(nbdContext) + } else { + err := run(ctx, nbdContext, *buildId, *mountPath, *verify) if err != nil { - return fmt.Errorf("failed to get object size: %w", err) - } - - h, err = header.NewHeader(&header.Metadata{ - BuildId: id, - BaseBuildId: id, - Size: uint64(size), - Version: 1, - BlockSize: header.RootfsBlockSize, - Generation: 1, - }, nil) - if err != nil { - return fmt.Errorf("failed to create header for rootfs without header: %w", err) + panic(fmt.Errorf("failed to mount rootfs: %w", err)) } } +} - diffCacheDir := filepath.Join(os.TempDir(), fmt.Sprintf("%s-rootfs.diff.cache-%s", buildID, uuid.New().String())) +func runEmpty(ctx, nbdContext context.Context, size int64, blockSize int64) error { + cowCachePath := filepath.Join(os.TempDir(), fmt.Sprintf("rootfs.ext4.cow.cache-%s", uuid.New().String())) - err = os.MkdirAll(diffCacheDir, 0o755) + emptyDevice, err := testutils.NewZeroDevice(size, blockSize) if err != nil { - return fmt.Errorf("failed to create diff cache directory: %w", err) + return fmt.Errorf("failed to create zero device: %w", err) } - defer os.RemoveAll(diffCacheDir) - - flags, err := featureflags.NewClient() - if err != nil { - return fmt.Errorf("failed to create feature flags client: %w", err) - } + defer os.RemoveAll(cowCachePath) - store, err := build.NewDiffStore( - nbdContext, - cfg.Config{}, - flags, - diffCacheDir, - 24*time.Hour, - 24*time.Hour, + cache, err := block.NewCache( + size, + blockSize, + cowCachePath, + false, ) if err != nil { - return fmt.Errorf("failed to create diff store: %w", err) + return fmt.Errorf("failed to create cache: %w", err) } - defer store.Close() + fmt.Printf("caching writes to: %+v\n", cowCachePath) - fmt.Printf("caching diffs to: %+v\n", diffCacheDir) + overlay := block.NewOverlay(emptyDevice, cache) + defer overlay.Close() - m, err := metrics.NewMetrics(noop.NewMeterProvider()) + devicePath, deviceCleanup, err := testutils.GetNBDDevice(nbdContext, testutils.NewLoggerOverlay(overlay)) + defer deviceCleanup.Run(ctx, 30*time.Second) if err != nil { - return fmt.Errorf("failed to create metrics: %w", err) + return fmt.Errorf("failed to get nbd device: %w", err) } - rootfs := build.NewFile(h, store, build.Rootfs, s, m) + fmt.Printf("rootfs exposed as device: %s\n", devicePath) + + <-ctx.Done() + + fmt.Println("closing rootfs mount") - readonlyDevice := newReadonlyDevice(rootfs, h, int64(h.Metadata.BlockSize)) + return nil +} + +func run(ctx, nbdContext context.Context, buildID, mountPath string, verify bool) error { + rootfs, rootfsCleanup, err := testutils.TemplateRootfs(ctx, buildID) + defer rootfsCleanup.Run(ctx, 30*time.Second) + if err != nil { + return fmt.Errorf("failed to get template rootfs: %w", err) + } cowCachePath := filepath.Join(os.TempDir(), fmt.Sprintf("%s-rootfs.ext4.cow.cache-%s", buildID, uuid.New().String())) defer os.RemoveAll(cowCachePath) cache, err := block.NewCache( - int64(h.Metadata.Size), - int64(h.Metadata.BlockSize), + int64(rootfs.Header().Metadata.Size), + int64(rootfs.Header().Metadata.BlockSize), cowCachePath, false, ) @@ -190,96 +133,72 @@ func mountRootfs(ctx, nbdContext context.Context, buildID, mountPath string) err fmt.Printf("caching writes to: %+v\n", cowCachePath) - overlay := block.NewOverlay(readonlyDevice, cache) + overlay := block.NewOverlay(rootfs, cache) defer overlay.Close() - devicePool, err := nbd.NewDevicePool() + devicePath, deviceCleanup, err := testutils.GetNBDDevice(nbdContext, overlay) + defer deviceCleanup.Run(ctx, 30*time.Second) if err != nil { - return fmt.Errorf("failed to create device pool: %w", err) + return fmt.Errorf("failed to get nbd device: %w", err) } - poolClosed := make(chan struct{}) - - defer func() { //nolint:contextcheck // we need to use separate context otherwise the cleanup can be problematic - <-poolClosed - - cleanupCtx, cancelCleanup := context.WithTimeout(context.Background(), 30*time.Second) - defer cancelCleanup() + fmt.Printf("rootfs exposed as device: %s\n", devicePath) - err = devicePool.Close(cleanupCtx) + if mountPath != "" { + err = os.MkdirAll(mountPath, 0o755) if err != nil { - fmt.Fprintf(os.Stderr, "failed to close device pool: %v\n", err) + return fmt.Errorf("failed to create mount path directory: %w", err) } - }() - poolCtx, poolCancel := context.WithCancel(nbdContext) - defer poolCancel() + fmt.Fprintf(os.Stdout, "creating mount path directory: %s\n", mountPath) - go func() { - devicePool.Populate(poolCtx) - close(poolClosed) - }() - - mnt := nbd.NewDirectPathMount(overlay, devicePool) - - mntIndex, err := mnt.Open(nbdContext) - if err != nil { - return fmt.Errorf("failed to open nbd mount: %w", err) - } - - defer func() { //nolint:contextcheck // we need to use separate context otherwise the cleanup can be problematic - cleanupCtx, cancelCleanup := context.WithTimeout(context.Background(), 30*time.Second) - defer cancelCleanup() - - err = mnt.Close(cleanupCtx) + mountCleanup, err := testutils.MountNBDDevice(devicePath, mountPath) + defer mountCleanup.Run(ctx, 30*time.Second) if err != nil { - fmt.Fprintf(os.Stderr, "failed to close nbd mount: %v\n", err) + return fmt.Errorf("failed to mount device to mount path: %w", err) } - }() - devicePath := nbd.GetDevicePath(mntIndex) - - fmt.Printf("rootfs exposed as device: %s\n", devicePath) + // We don't remove the dir as it might have been user created. - err = os.MkdirAll(mountPath, 0o755) - if err != nil { - return fmt.Errorf("failed to create mount path directory: %w", err) + fmt.Printf("rootfs mounted at path: %s\n", mountPath) } - fmt.Fprintf(os.Stdout, "creating mount path directory: %s\n", mountPath) + if verify { + fmt.Println("\nverifying rootfs integrity...") - // We don't remote the dir as it might have been user created. + cmd := exec.CommandContext(ctx, "e2fsck", "-nfv", devicePath) - err = unix.Mount(devicePath, mountPath, "ext4", unix.MS_RDONLY, "") - if err != nil { - return fmt.Errorf("failed to mount device to mount path: %w", err) - } + cmd.Stdout = os.Stdout + cmd.Stderr = os.Stderr - defer func() { - ticker := time.NewTicker(600 * time.Millisecond) - defer ticker.Stop() + err := cmd.Run() + if err != nil { + return fmt.Errorf("failed to verify rootfs integrity: %w", err) + } - cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 30*time.Second) - defer cleanupCancel() + fmt.Println("\nrootfs integrity verified") - for { - select { - case <-cleanupCtx.Done(): - fmt.Fprintf(os.Stderr, "failed to unmount device from mount path in time\n") + journalDir := filepath.Join(mountPath, "var", "log", "journal") + journalFiles, err := os.ReadDir(journalDir) + if err != nil { + return fmt.Errorf("failed to read journal directory: %w", err) + } - return - case <-ticker.C: - err = unix.Unmount(mountPath, 0) - if err == nil { - return - } + for _, journalFile := range journalFiles { + cmd := exec.CommandContext(ctx, "journalctl", "--verify", "--directory", filepath.Join(journalDir, journalFile.Name())) + cmd.Stdout = os.Stdout + cmd.Stderr = os.Stderr - fmt.Fprintf(os.Stderr, "failed to unmount device from mount path: %v\n", err) + err := cmd.Run() + if err != nil { + return fmt.Errorf("failed to verify journal file: %w", err) } } - }() - fmt.Printf("rootfs mounted at path: %s\n", mountPath) + fmt.Println("\njournal files verified") + + return nil + } <-ctx.Done() diff --git a/packages/orchestrator/cmd/mount-rootfs/start.sh b/packages/orchestrator/cmd/mount-rootfs/start.sh index 989f25f1bc..73da5567ab 100755 --- a/packages/orchestrator/cmd/mount-rootfs/start.sh +++ b/packages/orchestrator/cmd/mount-rootfs/start.sh @@ -2,13 +2,8 @@ set -euo pipefail -bucket=$1 -build=$2 -mount_path=$3 - go build -o ./bin/mount-rootfs ./cmd/mount-rootfs -TEMPLATE_BUCKET_NAME=$bucket \ GOOGLE_APPLICATION_CREDENTIALS="$HOME/.config/gcloud/application_default_credentials.json" \ sudo --preserve-env=GOOGLE_APPLICATION_CREDENTIALS,TEMPLATE_BUCKET_NAME \ -./bin/mount-rootfs -build $build -mount $mount_path +./bin/mount-rootfs "$@" \ No newline at end of file diff --git a/packages/orchestrator/go.mod b/packages/orchestrator/go.mod index a7dbb4a090..f7a6210bb0 100644 --- a/packages/orchestrator/go.mod +++ b/packages/orchestrator/go.mod @@ -41,7 +41,6 @@ require ( github.com/ngrok/firewall_toolkit v0.0.18 github.com/oapi-codegen/gin-middleware v1.0.2 github.com/pkg/errors v0.9.1 - github.com/pojntfx/go-nbd v0.3.2 github.com/shirou/gopsutil/v4 v4.25.6 github.com/soheilhy/cmux v0.1.5 github.com/stretchr/testify v1.11.1 diff --git a/packages/orchestrator/go.sum b/packages/orchestrator/go.sum index bf295884da..f7663be826 100644 --- a/packages/orchestrator/go.sum +++ b/packages/orchestrator/go.sum @@ -989,8 +989,6 @@ github.com/planetscale/vtprotobuf v0.6.1-0.20240319094008-0393e58bdf10/go.mod h1 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRIccs7FGNTlIRMkT8wgtp5eCXdBlqhYGL6U= github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= -github.com/pojntfx/go-nbd v0.3.2 h1:qI6S4qsHD87V9fTH6jiS4DIqq/rWmI0El0xSToMUDeg= -github.com/pojntfx/go-nbd v0.3.2/go.mod h1:SehHnbi2e8NiSAKby42Itm8SIoS7b+wAprsfPH3qgYk= github.com/posener/complete v1.1.1/go.mod h1:em0nMJCgc9GFtwrmVmEMR/ZL6WyhyjMBndrE9hABlRI= github.com/posener/complete v1.2.3/go.mod h1:WZIdtGGp+qx0sLrYKtIRAruyNpv6hFCicSgv7Sy7s/s= github.com/power-devops/perfstat v0.0.0-20240221224432-82ca36839d55 h1:o4JXh1EVt9k/+g42oCprj/FisM4qX9L3sZB3upGN2ZU= diff --git a/packages/orchestrator/internal/sandbox/build/cache.go b/packages/orchestrator/internal/sandbox/build/cache.go index 3706d03741..8d1016ccf9 100644 --- a/packages/orchestrator/internal/sandbox/build/cache.go +++ b/packages/orchestrator/internal/sandbox/build/cache.go @@ -303,3 +303,7 @@ func diskUsage(path string) (uint64, uint64, error) { return used, total, nil } + +func (s *DiffStore) RemoveCache() { + s.cache.DeleteAll() +} diff --git a/packages/orchestrator/internal/sandbox/nbd/path_direct_test.go b/packages/orchestrator/internal/sandbox/nbd/path_direct_test.go new file mode 100644 index 0000000000..8094985077 --- /dev/null +++ b/packages/orchestrator/internal/sandbox/nbd/path_direct_test.go @@ -0,0 +1,186 @@ +package nbd_test + +import ( + "context" + "crypto/rand" + "fmt" + "os" + "os/exec" + "path/filepath" + "testing" + "time" + + "github.com/google/uuid" + "github.com/stretchr/testify/require" + "golang.org/x/sys/unix" + + "github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/block" + "github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/nbd/testutils" + "github.com/e2b-dev/infra/packages/shared/pkg/storage/header" +) + +func TestPathDirect_Direct4MBWrite(t *testing.T) { + size := int64(10 * 1024 * 1024) + + deviceFile := setupNBDDevice(t, size, header.RootfsBlockSize, unix.O_DIRECT|unix.O_RDWR) + + const bs = 4 * 1024 * 1024 + buf, err := unix.Mmap(-1, 0, bs, unix.PROT_READ|unix.PROT_WRITE, unix.MAP_SHARED|unix.MAP_ANON) + if err != nil { + panic(err) + } + + t.Cleanup(func() { + unix.Munmap(buf) + }) + + n, err := deviceFile.WriteAt(buf, 0) + require.NoError(t, err, "failed to write to device") + require.Equal(t, len(buf), n, "partial write") + + readData := make([]byte, bs) + n, err = deviceFile.ReadAt(readData, 0) + require.NoError(t, err, "failed to read from device") + require.Equal(t, len(readData), n, "partial read") + require.Equal(t, buf, readData, "data mismatch") +} + +// We usually see the 32MB write be split into smaller writes, even on O_DIRECT. +func TestPathDirect_Direct32MBWrite(t *testing.T) { + size := int64(256 * 1024 * 1024) + + deviceFile := setupNBDDevice(t, size, header.RootfsBlockSize, unix.O_DIRECT|unix.O_RDWR) + + const bs = 32 * 1024 * 1024 + buf, err := unix.Mmap(-1, 0, bs, unix.PROT_READ|unix.PROT_WRITE, unix.MAP_SHARED|unix.MAP_ANON) + require.NoError(t, err, "failed to mmap") + + t.Cleanup(func() { + unix.Munmap(buf) + }) + + n, err := deviceFile.WriteAt(buf, 0) + require.NoError(t, err, "failed to write to device") + require.Equal(t, len(buf), n, "partial write") + + readData := make([]byte, bs) + n, err = deviceFile.ReadAt(readData, 0) + require.NoError(t, err, "failed to read from device") + require.Equal(t, len(readData), n, "partial read") + require.Equal(t, buf, readData, "data mismatch") +} + +func TestPathDirect_Write(t *testing.T) { + size := int64(5 * 1024 * 1024) + + deviceFile := setupNBDDevice(t, size, header.RootfsBlockSize, os.O_RDWR) + + const writeSize = 1024 * 1024 + testData := make([]byte, writeSize) + _, err := rand.Read(testData) + require.NoError(t, err, "failed to generate random data") + + n, err := deviceFile.WriteAt(testData, 0) + require.NoError(t, err, "failed to write data to device") + require.Equal(t, len(testData), n, "partial write") + + readData := make([]byte, writeSize) + n, err = deviceFile.ReadAt(readData, 0) + require.NoError(t, err, "failed to read data from device") + require.Equal(t, len(readData), n, "partial read") + require.Equal(t, testData, readData, "data mismatch") +} + +func TestPathDirect_WriteAtOffset(t *testing.T) { + size := int64(5 * 1024 * 1024) + + deviceFile := setupNBDDevice(t, size, header.RootfsBlockSize, os.O_RDWR) + + const writeSize = 512 * 1024 + const writeOffset = 512 * 1024 + testData := make([]byte, writeSize) + _, err := rand.Read(testData) + require.NoError(t, err, "failed to generate random data") + + n, err := deviceFile.WriteAt(testData, writeOffset) + require.NoError(t, err, "failed to write data to device") + require.Equal(t, len(testData), n, "partial write") + + readData := make([]byte, writeSize) + n, err = deviceFile.ReadAt(readData, writeOffset) + require.NoError(t, err, "failed to read data from device") + require.Equal(t, len(readData), n, "partial read") + require.Equal(t, testData, readData, "data mismatch") +} + +func TestPathDirect_LargeWrite(t *testing.T) { + size := int64(1200 * 1024 * 1024) + + deviceFile := setupNBDDevice(t, size, header.RootfsBlockSize, os.O_RDWR) + + time.Sleep(1 * time.Second) + cmd := exec.CommandContext(t.Context(), "dd", "if=/dev/zero", "of="+deviceFile.Name(), "bs=1G", "count=1") + cmd.Stdout = os.Stdout + cmd.Stderr = os.Stderr + + err := cmd.Run() + require.NoError(t, err, "failed to execute dd command") +} + +func TestPathLargeRead(t *testing.T) { + size := int64(1200 * 1024 * 1024) + + deviceFile := setupNBDDevice(t, size, header.RootfsBlockSize, os.O_RDONLY) + time.Sleep(1 * time.Second) + + cmd := exec.CommandContext(t.Context(), "dd", "if="+deviceFile.Name(), "of=/dev/null", "bs=1G", "count=1") + cmd.Stdout = os.Stdout + cmd.Stderr = os.Stderr + + err := cmd.Run() + require.NoError(t, err, "failed to execute dd command") +} + +func setupNBDDevice(t *testing.T, size, blockSize int64, flags int) *os.File { + t.Helper() + + require.Equal(t, 0, os.Geteuid(), "the nbd requires root privileges to run") + + emptyDevice, err := testutils.NewZeroDevice(size, blockSize) + require.NoError(t, err, "failed to create zero device") + + cowCachePath := filepath.Join(os.TempDir(), fmt.Sprintf("test-rootfs.ext4.cow.cache-%s", uuid.New().String())) + t.Cleanup(func() { + os.RemoveAll(cowCachePath) + }) + + cache, err := block.NewCache( + size, + blockSize, + cowCachePath, + false, + ) + require.NoError(t, err, "failed to create cache") + + overlay := block.NewOverlay(emptyDevice, cache) + t.Cleanup(func() { + overlay.Close() + }) + + nbdContext := context.Background() + devicePath, deviceCleanup, err := testutils.GetNBDDevice(nbdContext, overlay) + t.Cleanup(func() { + deviceCleanup.Run(t.Context(), 30*time.Second) + }) + require.NoError(t, err, "failed to get nbd device") + + t.Logf("NBD device path: %s", devicePath) + + deviceFile, err := os.OpenFile(devicePath, flags, 0) + require.NoError(t, err, "failed to open device") + t.Cleanup(func() { + deviceFile.Close() + }) + + return deviceFile +} diff --git a/packages/orchestrator/internal/sandbox/nbd/testutils/build_device.go b/packages/orchestrator/internal/sandbox/nbd/testutils/build_device.go new file mode 100644 index 0000000000..4f167834e1 --- /dev/null +++ b/packages/orchestrator/internal/sandbox/nbd/testutils/build_device.go @@ -0,0 +1,40 @@ +package testutils + +import ( + "github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/block" + "github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/build" + "github.com/e2b-dev/infra/packages/shared/pkg/storage/header" +) + +var _ block.ReadonlyDevice = (*BuildDevice)(nil) + +type BuildDevice struct { + *build.File + + header *header.Header + blockSize int64 +} + +func NewBuildDevice(file *build.File, header *header.Header, blockSize int64) *BuildDevice { + return &BuildDevice{ + File: file, + header: header, + blockSize: blockSize, + } +} + +func (m *BuildDevice) Close() error { + return nil +} + +func (m *BuildDevice) BlockSize() int64 { + return m.blockSize +} + +func (m *BuildDevice) Header() *header.Header { + return m.header +} + +func (m *BuildDevice) Size() (int64, error) { + return int64(m.header.Metadata.Size), nil +} diff --git a/packages/orchestrator/internal/sandbox/nbd/testutils/cleaner.go b/packages/orchestrator/internal/sandbox/nbd/testutils/cleaner.go new file mode 100644 index 0000000000..59cf3ee7a6 --- /dev/null +++ b/packages/orchestrator/internal/sandbox/nbd/testutils/cleaner.go @@ -0,0 +1,39 @@ +package testutils + +import ( + "context" + "errors" + "fmt" + "slices" + "sync" + "time" +) + +type Cleaner struct { + steps []func(ctx context.Context) error + once sync.Once +} + +func (c *Cleaner) Add(f func(ctx context.Context) error) { + c.steps = append(c.steps, f) +} + +func (c *Cleaner) Run(ctx context.Context, timeout time.Duration) (err error) { + ctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), timeout) + defer cancel() + + c.once.Do(func() { + var errs []error + + for _, step := range slices.Backward(c.steps) { + err := step(ctx) + if err != nil { + errs = append(errs, fmt.Errorf("failed to run step: %w", err)) + } + } + + err = errors.Join(errs...) + }) + + return err +} diff --git a/packages/orchestrator/internal/sandbox/nbd/testutils/logger_overlay.go b/packages/orchestrator/internal/sandbox/nbd/testutils/logger_overlay.go new file mode 100644 index 0000000000..8eb5a0f65a --- /dev/null +++ b/packages/orchestrator/internal/sandbox/nbd/testutils/logger_overlay.go @@ -0,0 +1,74 @@ +package testutils + +import ( + "context" + "fmt" + "os" + + "github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/block" + "github.com/e2b-dev/infra/packages/shared/pkg/storage/header" +) + +type LoggerOverlay struct { + overlay *block.Overlay +} + +func NewLoggerOverlay(overlay *block.Overlay) *LoggerOverlay { + return &LoggerOverlay{overlay: overlay} +} + +func (l *LoggerOverlay) ReadAt(ctx context.Context, p []byte, off int64) (int, error) { + defer func() { + if r := recover(); r != nil { + fmt.Fprintf(os.Stdout, "[read panic recovered]: [%d, %d] -> %v\n", off, len(p), r) + } + }() + + fmt.Fprintf(os.Stdout, "[read started]: [%d, %d]\n", off, len(p)) + + n, err := l.overlay.ReadAt(ctx, p, off) + + fmt.Fprintf(os.Stdout, "[read completed]: [%d, %d] -> %d\n", off, len(p), n) + + return n, err +} + +func (l *LoggerOverlay) WriteAt(p []byte, off int64) (int, error) { + defer func() { + if r := recover(); r != nil { + fmt.Fprintf(os.Stdout, "[write panic recovered]: [%d, %d] -> %v\n", off, len(p), r) + } + }() + + fmt.Fprintf(os.Stdout, "[write started]: [%d, %d]\n", off, len(p)) + + n, err := l.overlay.WriteAt(p, off) + + fmt.Fprintf(os.Stdout, "[write completed]: [%d, %d] -> %d\n", off, len(p), n) + + return n, err +} + +func (l *LoggerOverlay) Size() (int64, error) { + return l.overlay.Size() +} + +func (l *LoggerOverlay) BlockSize() int64 { + return l.overlay.BlockSize() +} + +func (l *LoggerOverlay) Header() *header.Header { + return l.overlay.Header() +} + +func (l *LoggerOverlay) Close() error { + return l.overlay.Close() +} + +func (l *LoggerOverlay) EjectCache() (*block.Cache, error) { + return l.overlay.EjectCache() +} + +func (l *LoggerOverlay) Slice(ctx context.Context, off, length int64) ([]byte, error) { + return l.overlay.Slice(ctx, off, length) +} diff --git a/packages/orchestrator/internal/sandbox/nbd/testutils/mount.go b/packages/orchestrator/internal/sandbox/nbd/testutils/mount.go new file mode 100644 index 0000000000..cab0033e8c --- /dev/null +++ b/packages/orchestrator/internal/sandbox/nbd/testutils/mount.go @@ -0,0 +1,44 @@ +package testutils + +import ( + "context" + "fmt" + "os" + "time" + + "golang.org/x/sys/unix" + + "github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/nbd" +) + +func MountNBDDevice(device nbd.DevicePath, mountPath string) (*Cleaner, error) { + var cleaner Cleaner + + err := unix.Mount(device, mountPath, "ext4", 0, "") + if err != nil { + return &cleaner, fmt.Errorf("failed to mount device to mount path: %w", err) + } + + cleaner.Add(func(cleanupCtx context.Context) error { + ticker := time.NewTicker(600 * time.Millisecond) + defer ticker.Stop() + + for { + select { + case <-cleanupCtx.Done(): + fmt.Fprintf(os.Stderr, "failed to unmount device from mount path in time\n") + + return nil + case <-ticker.C: + err = unix.Unmount(mountPath, 0) + if err == nil { + return nil + } + + fmt.Fprintf(os.Stderr, "failed to unmount device from mount path: %v\n", err) + } + } + }) + + return &cleaner, nil +} diff --git a/packages/orchestrator/internal/sandbox/nbd/testutils/nbd_device.go b/packages/orchestrator/internal/sandbox/nbd/testutils/nbd_device.go new file mode 100644 index 0000000000..1b53d25aa4 --- /dev/null +++ b/packages/orchestrator/internal/sandbox/nbd/testutils/nbd_device.go @@ -0,0 +1,63 @@ +package testutils + +import ( + "context" + "fmt" + "os" + + "github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/block" + "github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/nbd" +) + +func GetNBDDevice(ctx context.Context, backend block.Device) (nbd.DevicePath, *Cleaner, error) { + var cleaner Cleaner + + devicePool, err := nbd.NewDevicePool() + if err != nil { + return "", &cleaner, fmt.Errorf("failed to create device pool: %w", err) + } + + poolClosed := make(chan struct{}) + + cleaner.Add(func(cleanupCtx context.Context) error { + <-poolClosed + + err = devicePool.Close(cleanupCtx) + if err != nil { + fmt.Fprintf(os.Stderr, "failed to close device pool: %v\n", err) + } + + return nil + }) + + poolCtx, poolCancel := context.WithCancel(ctx) + + cleaner.Add(func(context.Context) error { + poolCancel() + + return nil + }) + + go func() { + devicePool.Populate(poolCtx) + close(poolClosed) + }() + + mnt := nbd.NewDirectPathMount(backend, devicePool) + + mntIndex, err := mnt.Open(ctx) + if err != nil { + return "", &cleaner, fmt.Errorf("failed to open nbd mount: %w", err) + } + + cleaner.Add(func(cleanupCtx context.Context) error { + err = mnt.Close(cleanupCtx) + if err != nil { + fmt.Fprintf(os.Stderr, "failed to close nbd mount: %v\n", err) + } + + return nil + }) + + return nbd.GetDevicePath(mntIndex), &cleaner, nil +} diff --git a/packages/orchestrator/internal/sandbox/nbd/testutils/template_rootfs.go b/packages/orchestrator/internal/sandbox/nbd/testutils/template_rootfs.go new file mode 100644 index 0000000000..b039c7679c --- /dev/null +++ b/packages/orchestrator/internal/sandbox/nbd/testutils/template_rootfs.go @@ -0,0 +1,120 @@ +package testutils + +import ( + "context" + "fmt" + "os" + "path/filepath" + "time" + + "github.com/google/uuid" + "go.opentelemetry.io/otel/metric/noop" + + "github.com/e2b-dev/infra/packages/orchestrator/internal/cfg" + "github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/block/metrics" + "github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/build" + featureflags "github.com/e2b-dev/infra/packages/shared/pkg/feature-flags" + "github.com/e2b-dev/infra/packages/shared/pkg/storage" + "github.com/e2b-dev/infra/packages/shared/pkg/storage/header" +) + +func TemplateRootfs(ctx context.Context, buildID string) (*BuildDevice, *Cleaner, error) { + var cleaner Cleaner + + files := storage.TemplateFiles{ + BuildID: buildID, + } + + s, err := storage.GetTemplateStorageProvider(ctx, nil) + if err != nil { + return nil, &cleaner, fmt.Errorf("failed to get storage provider: %w", err) + } + + obj, err := s.OpenObject(ctx, files.StorageRootfsHeaderPath(), storage.RootFSHeaderObjectType) + if err != nil { + return nil, &cleaner, fmt.Errorf("failed to open object: %w", err) + } + + h, err := header.Deserialize(ctx, obj) + if err != nil { + id, err := uuid.Parse(buildID) + if err != nil { + return nil, &cleaner, fmt.Errorf("failed to parse build id: %w", err) + } + + r, err := s.OpenSeekableObject(ctx, files.StorageRootfsPath(), storage.RootFSObjectType) + if err != nil { + return nil, &cleaner, fmt.Errorf("failed to open object: %w", err) + } + + size, err := r.Size(ctx) + if err != nil { + return nil, &cleaner, fmt.Errorf("failed to get object size: %w", err) + } + + h, err = header.NewHeader(&header.Metadata{ + BuildId: id, + BaseBuildId: id, + Size: uint64(size), + Version: 1, + BlockSize: header.RootfsBlockSize, + Generation: 1, + }, nil) + if err != nil { + return nil, &cleaner, fmt.Errorf("failed to create header for rootfs without header: %w", err) + } + } + + diffCacheDir := filepath.Join(os.TempDir(), fmt.Sprintf("%s-rootfs.diff.cache-%s", buildID, uuid.New().String())) + + err = os.MkdirAll(diffCacheDir, 0o755) + if err != nil { + return nil, &cleaner, fmt.Errorf("failed to create diff cache directory: %w", err) + } + + cleaner.Add(func(context.Context) error { + return os.RemoveAll(diffCacheDir) + }) + + flags, err := featureflags.NewClient() + if err != nil { + return nil, &cleaner, fmt.Errorf("failed to create feature flags client: %w", err) + } + + store, err := build.NewDiffStore( + ctx, + cfg.Config{}, + flags, + diffCacheDir, + 24*time.Hour, + 24*time.Hour, + ) + if err != nil { + return nil, &cleaner, fmt.Errorf("failed to create diff store: %w", err) + } + + cleaner.Add(func(context.Context) error { + store.RemoveCache() + + return nil + }) + + cleaner.Add(func(context.Context) error { + store.Close() + + return nil + }) + + m, err := metrics.NewMetrics(noop.NewMeterProvider()) + if err != nil { + return nil, &cleaner, fmt.Errorf("failed to create metrics: %w", err) + } + + buildDevice := NewBuildDevice( + build.NewFile(h, store, build.Rootfs, s, m), + h, + int64(h.Metadata.BlockSize), + ) + + return buildDevice, &cleaner, nil +} diff --git a/packages/orchestrator/internal/sandbox/nbd/testutils/zero_device.go b/packages/orchestrator/internal/sandbox/nbd/testutils/zero_device.go new file mode 100644 index 0000000000..f266a1a2fb --- /dev/null +++ b/packages/orchestrator/internal/sandbox/nbd/testutils/zero_device.go @@ -0,0 +1,71 @@ +package testutils + +import ( + "context" + "fmt" + + "github.com/google/uuid" + + "github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/block" + "github.com/e2b-dev/infra/packages/shared/pkg/storage/header" +) + +var _ block.ReadonlyDevice = (*ZeroDevice)(nil) + +type ZeroDevice struct { + blockSize int64 + size int64 + header *header.Header +} + +func NewZeroDevice(size int64, blockSize int64) (*ZeroDevice, error) { + h, err := header.NewHeader(header.NewTemplateMetadata( + uuid.Nil, + uint64(blockSize), + uint64(size), + ), + []*header.BuildMap{ + { + Offset: 0, + Length: uint64(size), + BuildId: uuid.Nil, + BuildStorageOffset: 0, + }, + }, + ) + if err != nil { + return nil, fmt.Errorf("failed to create header: %w", err) + } + + return &ZeroDevice{ + size: size, + blockSize: blockSize, + header: h, + }, nil +} + +func (z *ZeroDevice) ReadAt(_ context.Context, p []byte, _ int64) (n int, err error) { + clear(p) + + return len(p), nil +} + +func (z *ZeroDevice) BlockSize() int64 { + return z.blockSize +} + +func (z *ZeroDevice) Slice(_ context.Context, _, length int64) ([]byte, error) { + return make([]byte, length), nil +} + +func (z *ZeroDevice) Header() *header.Header { + return z.header +} + +func (z *ZeroDevice) Close() error { + return nil +} + +func (z *ZeroDevice) Size() (int64, error) { + return z.size, nil +} diff --git a/packages/shared/pkg/storage/header/header.go b/packages/shared/pkg/storage/header/header.go index 74b51be985..fd9e4dfc34 100644 --- a/packages/shared/pkg/storage/header/header.go +++ b/packages/shared/pkg/storage/header/header.go @@ -85,6 +85,7 @@ func (t *Header) GetShiftedMapping(offset int64) (mappedOffset int64, mappedLeng return mappedOffset, mappedLength, buildID, nil } +// TODO: Maybe we can optimize mapping by automatically assuming the mapping is uuid.Nil if we don't find it + stopping storing the nil mapping. func (t *Header) getMapping(offset int64) (*BuildMap, int64, error) { if offset < 0 || offset >= int64(t.Metadata.Size) { if t.IsNormalizeFixApplied() { diff --git a/packages/shared/pkg/storage/storage_cache_test.go b/packages/shared/pkg/storage/storage_cache_test.go index fa5950c1c2..86f89243b3 100644 --- a/packages/shared/pkg/storage/storage_cache_test.go +++ b/packages/shared/pkg/storage/storage_cache_test.go @@ -236,25 +236,3 @@ func TestMoveWithoutReplace_FailWhenExists(t *testing.T) { _, err = os.Stat(src) assert.ErrorIs(t, err, os.ErrNotExist) } - -func TestMoveWithoutReplace_Fail(t *testing.T) { - td := t.TempDir() - content := []byte("alpha") - src := filepath.Join(td, "src") - require.NoError(t, os.WriteFile(src, content, 0o644)) - - roDir := filepath.Join(td, "ro") - require.NoError(t, os.Mkdir(roDir, 0o555)) // r-x only, no write - t.Cleanup(func() { - // ensure cleanup possible - err := os.Chmod(roDir, 0o755) - assert.NoError(t, err) - }) - - dst := filepath.Join(roDir, "dst") - err := moveWithoutReplace(src, dst) - require.Error(t, err) - - _, err = os.Stat(src) - assert.ErrorIs(t, err, os.ErrNotExist) -}