package copy_on_write_test

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"math/rand"
	"os"
	"path/filepath"
	"runtime"
	"strconv"
	"sync"
	"syscall"
	"testing"

	"github.com/buildbuddy-io/buildbuddy/enterprise/server/remote_execution/copy_on_write"
	"github.com/buildbuddy-io/buildbuddy/enterprise/server/remote_execution/copy_on_write/cow_cgo_testutil"
	"github.com/buildbuddy-io/buildbuddy/enterprise/server/remote_execution/snaputil"
	"github.com/buildbuddy-io/buildbuddy/server/interfaces"
	"github.com/buildbuddy-io/buildbuddy/server/metrics"
	"github.com/buildbuddy-io/buildbuddy/server/remote_cache/digest"
	"github.com/buildbuddy-io/buildbuddy/server/resources"
	"github.com/buildbuddy-io/buildbuddy/server/testutil/testenv"
	"github.com/buildbuddy-io/buildbuddy/server/testutil/testfs"
	"github.com/buildbuddy-io/buildbuddy/server/testutil/testmetrics"
	"github.com/buildbuddy-io/buildbuddy/server/util/disk"
	"github.com/buildbuddy-io/buildbuddy/server/util/log"
	"github.com/buildbuddy-io/buildbuddy/server/util/testing/flags"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/stretchr/testify/require"
	"golang.org/x/sync/errgroup"

	repb "github.com/buildbuddy-io/buildbuddy/proto/remote_execution"
)

const (
	backingFileSizeBytes int64 = 32 * 1024
)

func init() {
	// Ensure some memory is allocated for the shared LRU.
	if err := resources.Configure(true /*=snapshotSharingEnabled*/); err != nil {
		log.Fatalf("Failed to configure resources: %s", err)
	}
}

func TestMmap(t *testing.T) {
	s, path := newMmap(t)
	testStore(t, s, path)
}

func TestMmap_Digest(t *testing.T) {
	s, _ := newMmap(t)
	actualDigest1, err := s.Digest()
	require.NoError(t, err)
	r, err := interfaces.StoreReader(s)
	require.NoError(t, err)
	expectedDigest, err := digest.Compute(r, repb.DigestFunction_BLAKE3)
	require.NoError(t, err)
	require.Equal(t, expectedDigest, actualDigest1)

	// Write random bytes and test digest again
	randomBuf := randBytes(t, 1024)
	n, err := s.WriteAt(randomBuf, 100)
	require.NoError(t, err)
	require.Equal(t, len(randomBuf), n)
	actualDigest2, err := s.Digest()
	require.NoError(t, err)
	r, err = interfaces.StoreReader(s)
	require.NoError(t, err)
	expectedDigest, err = digest.Compute(r, repb.DigestFunction_BLAKE3)
	require.NoError(t, err)
	require.Equal(t, expectedDigest, actualDigest2)
	require.NotEqual(t, actualDigest1, actualDigest2)

}

func TestMmap_Concurrency(t *testing.T) {
	s, _ := newMmap(t)

	eg := &errgroup.Group{}
	for i := 0; i < 100; i++ {
		eg.Go(func() error {
			s.Source()
			return nil
		})
		eg.Go(func() error {
			buf := make([]byte, 100)
			_, err := rand.Read(buf)
			if err != nil {
				return err
			}
			_, err = s.WriteAt(buf, 0)
			return err
		})
		eg.Go(func() error {
			buf := make([]byte, 100)
			_, err := s.ReadAt(buf, 0)
			return err
		})
		eg.Go(func() error {
			return s.Sync()
		})
		eg.Go(func() error {
			_, err := s.SizeBytes()
			return err
		})
		eg.Go(func() error {
			_, err := s.StartAddress()
			return err
		})
		eg.Go(func() error {
			return s.Fetch()
		})
		eg.Go(func() error {
			// Should be safe to call Unmap at any time. ReadAt / WriteAt should
			// re-map if needed.
			return s.Unmap()
		})
	}
	err := eg.Wait()
	require.NoError(t, err)
	err = s.Close()
	require.NoError(t, err)
}

func TestCOW_Basic(t *testing.T) {
	ctx := context.Background()
	env := testenv.GetTestEnv(t)
	path := makeEmptyTempFile(t, backingFileSizeBytes)
	dataDir := testfs.MakeTempDir(t)
	chunkSizeBytes := backingFileSizeBytes / 2
	s, err := copy_on_write.ConvertFileToCOW(ctx, env, path, chunkSizeBytes, dataDir, "", false)
	require.NoError(t, err)
	// Don't validate against the backing file, since COWFromFile makes a copy
	// of the underlying file.
	testStore(t, s, "" /*=path*/)
}

func TestCOW_Concurrency(t *testing.T) {
	const chunkSizeBytes = 1024 * 512
	const fileSizeBytes = chunkSizeBytes * 20

	ctx := context.Background()
	env := testenv.GetTestEnv(t)
	path := makeEmptyTempFile(t, backingFileSizeBytes)
	dataDir := testfs.MakeTempDir(t)
	s, err := copy_on_write.ConvertFileToCOW(ctx, env, path, chunkSizeBytes, dataDir, "", false)
	require.NoError(t, err)

	tester := NewStoreTester(t, s)

	eg := &errgroup.Group{}
	eg.SetLimit(1000)
	for i := 0; i < 10_000; i++ {
		eg.Go(func() error {
			if rand.Float64() < 0.2 {
				tester.WriteRandomRange()
			} else {
				tester.ReadRandomRange()
			}
			return nil
		})
	}

	// Note: Sync will likely be called while the above operations are
	// still in progress, which should be OK.

	err = s.Sync()
	require.NoError(t, err)

	eg.Wait()

	tester.ReadAll()

	err = s.Close()
	require.NoError(t, err)
}

func TestCOW_SparseData(t *testing.T) {
	if runtime.GOOS == "darwin" {
		// Sparse files work a bit differently on macOS, just skip for now.
		t.SkipNow()
	}

	ctx := context.Background()
	env := testenv.GetTestEnv(t)
	// Figure out the IO block size (number of bytes transferred to/from disk
	// for each IO operation). This is the minimum seek size when using seek()
	// with SEEK_DATA.
	tmp := testfs.MakeTempDir(t)
	ioBlockSize := ioBlockSize(t, tmp)
	// Use a chunk size that is a few times larger than the IO block size.
	const blocksPerChunk = 4
	chunkSize := ioBlockSize * blocksPerChunk
	// Write a file consisting of the following chunks:
	// - 0: completely empty
	// - 1: data block at the beginning of the chunk
	// - 2: data block somewhere in the middle of the chunk
	// - 3: data block at the end of the chunk
	chunks := make([][]byte, 4)
	for i := 0; i < len(chunks); i++ {
		chunks[i] = make([]byte, chunkSize)
	}
	// chunkData[0]: empty
	chunks[1][0] = 1
	chunks[2][ioBlockSize] = 1
	chunks[3][chunkSize-1] = 1
	data := concatBytes(chunks...)
	dataFilePath := filepath.Join(tmp, "data.bin")
	writeSparseFile(t, dataFilePath, data, ioBlockSize)
	outDir := testfs.MakeTempDir(t)

	// Now split the file.
	c, err := copy_on_write.ConvertFileToCOW(ctx, env, dataFilePath, chunkSize, outDir, "", false)
	require.NoError(t, err)
	t.Cleanup(func() { c.Close() })

	// Make sure the full content matches our buffer.
	dataOut := make([]byte, len(data))
	n, err := c.ReadAt(dataOut, 0)
	require.NoError(t, err)
	require.Equal(t, len(dataOut), n)

	// Inspect the chunk files and ensure they have the expected physical size.
	for i := 0; i < len(chunks); i++ {
		chunkPath := filepath.Join(outDir, strconv.Itoa(i*int(chunkSize)))
		// We wrote one data block per chunk except for the one chunk that was
		// all empty. The empty chunk should not have written a file.
		if i == 0 {
			exists, err := disk.FileExists(context.TODO(), chunkPath)
			require.NoError(t, err)
			require.False(t, exists, "chunk 0 should not exist on disk")
			continue
		}
		nb := numIOBlocks(t, chunkPath)
		require.Equal(t, int64(1), nb, "chunk %d IO block count", i)
	}

	// Now write a single data byte to the empty chunk (0), and sync it to disk.
	n, err = c.WriteAt([]byte{1}, 0)
	require.NoError(t, err)
	require.Equal(t, 1, n)
	err = c.SortedChunks()[0].Sync()
	require.NoError(t, err)

	// Make sure we wrote a dirty chunk with the expected contents, and only a
	// single IO block on disk (since we only wrote 1 byte).
	dirtyPath := filepath.Join(outDir, "0.dirty")
	b, err := os.ReadFile(dirtyPath)
	require.NoError(t, err)
	expectedContent := make([]byte, len(chunks[0]))
	expectedContent[0] = 1
	require.Equal(t, expectedContent, b)

	// TODO(http://go/b/3918): figure out why mmap writes create 4 physical
	// blocks on baremetal executors, and assert that only 1 physical block is
	// written (if possible).
	ioBlocksInDirtyFile := numIOBlocks(t, dirtyPath)
	require.True(t, ioBlocksInDirtyFile == 1 || ioBlocksInDirtyFile == 4, "unexpected number of IO blocks (%d)", ioBlocksInDirtyFile)
}

func TestCOW_Resize(t *testing.T) {
	const chunkSize = 2 * 4096
	ctx := context.Background()
	env := testenv.GetTestEnv(t)
	for _, test := range []struct {
		Name             string
		OldSize, NewSize int64
		ExpectError      bool
	}{
		{Name: "AddNewChunkOnly", OldSize: chunkSize, NewSize: chunkSize + 1},
		{Name: "RightPadLastChunk", OldSize: chunkSize + 1, NewSize: chunkSize + 2},
		{Name: "RightPadLastChunkAndAddNewChunk", OldSize: chunkSize + 1, NewSize: chunkSize * 3},
		{Name: "RightPadLastChunkAndAddMultipleNewChunks", OldSize: chunkSize + 1, NewSize: chunkSize * 4},
		{Name: "DecreaseSize", OldSize: chunkSize, NewSize: chunkSize - 1, ExpectError: true},
	} {
		t.Run(test.Name, func(t *testing.T) {
			for i := 0; i < 10; i++ {
				// Start out with a file containing random data
				startBuf := randBytes(t, int(test.OldSize))
				src := makeTempFile(t, startBuf)
				dir := testfs.MakeTempDir(t)
				cow, err := copy_on_write.ConvertFileToCOW(ctx, env, src, chunkSize, dir, "", false)
				require.NoError(t, err)

				// Resize the COW
				oldSize, err := cow.Resize(test.NewSize)
				if test.ExpectError {
					require.Error(t, err)
					continue
				}
				require.NoError(t, err)
				require.Equal(t, test.OldSize, oldSize)

				// Read random ranges; should match startBuf right-padded with
				// zeroes.
				startRightPad := append(startBuf, make([]byte, test.NewSize-test.OldSize)...)
				for i := 0; i < 10; i++ {
					offset, length := randSubslice(int(test.NewSize))
					b := make([]byte, length)
					_, err := cow.ReadAt(b, int64(offset))
					require.NoError(t, err)
					for j, b := range b {
						require.Equal(t, startRightPad[offset+j], b)
					}
				}

				endBuf := make([]byte, test.NewSize)
				copy(endBuf, startBuf)
				// Fill a random data range in the resized file
				offset, length := randSubslice(int(test.NewSize))
				copy(endBuf[offset:offset+length], randBytes(t, length))
				_, err = cow.WriteAt(endBuf[offset:offset+length], int64(offset))
				require.NoError(t, err)
				// Now read back the whole COW and make sure it matches our
				// expected data.
				cowReader, err := interfaces.StoreReader(cow)
				require.NoError(t, err)
				b, err := io.ReadAll(cowReader)
				require.NoError(t, err)
				require.True(t, bytes.Equal(endBuf, b))

				// Read random ranges again; should match endBuf this time.
				for i := 0; i < 10; i++ {
					offset, length := randSubslice(int(test.NewSize))
					b := make([]byte, length)
					_, err := cow.ReadAt(b, int64(offset))
					require.NoError(t, err)
					for j, b := range b {
						require.Equal(t, endBuf[offset+j], b)
					}
				}
			}
		})
	}
}

func TestCOW_MmapLRUDoesNotDeadlock(t *testing.T) {
	// Set a relatively small LRU size limit to increase LRU contention.
	flags.Set(t, "executor.mmap_memory_bytes", 64*1024*1024)
	err := resources.Configure(true /*=enableSnapshotSharing*/)
	require.NoError(t, err)
	copy_on_write.ResetSharedLRUForTest()
	copy_on_write.ResetMmmapedBytesMetricForTest()

	ctx := context.Background()
	env := testenv.GetTestEnv(t)

	const fileSize = 200 * 1024 * 1024
	const chunkSize = 4000 * 1024

	f := testfs.CreateTemp(t)
	err = f.Truncate(fileSize)
	require.NoError(t, err)
	path := f.Name()
	err = f.Close()
	require.NoError(t, err)

	chunkDir := testfs.MakeTempDir(t)
	cow, err := copy_on_write.ConvertFileToCOW(ctx, env, path, chunkSize, chunkDir, "", false)
	require.NoError(t, err)

	var eg errgroup.Group
	eg.SetLimit(100)
	for i := 0; i < 10_000; i++ {
		eg.Go(func() error {
			p := make([]byte, 1)
			offset := rand.Int63n(fileSize - 1)
			r := rand.Float64()
			if r < 0.02 {
				return cow.UnmapChunk(offset)
			} else if r < 0.04 {
				_, err := cow.WriteAt(p, offset)
				return err
			} else {
				_, err := cow.ReadAt(p, offset)
				return err
			}
		})
	}
	err = eg.Wait()
	require.NoError(t, err)

	// Unmap everything again and assert that the LRU accounting is correct.
	err = cow.Close()
	require.NoError(t, err)

	n := testmetrics.GaugeValueForLabels(
		t, metrics.COWSnapshotMemoryMappedBytes,
		prometheus.Labels{metrics.FileName: filepath.Base(chunkDir)})
	require.Equal(t, float64(0), n)
}

func BenchmarkCOW_ReadWritePerformance(b *testing.B) {
	flags.Set(b, "app.log_level", "error")
	log.Configure()

	const chunkSize = 4000 * 1024 // 4MB
	// TODO: figure out a more realistic distribution of read/write size
	const ioSize = 4096
	// Use a relatively small disk size to avoid expensive test setup.
	const diskSize = 64 * 1024 * 1024
	// Read/write a total volume equal to the disk size so that sequential
	// read/write tests don't touch the same block twice.
	const ioCountPerBenchOp = diskSize / ioSize

	ctx := context.Background()
	env := testenv.GetTestEnv(b)

	for _, test := range []struct {
		name string

		// Value 0-1 indicating the approx fraction of data blocks in the
		// initial file (before chunking). Non-data blocks will be empty.
		initialDensity float64

		// Value 0-1 indicating what fraction of requests are reads.
		readFraction float64
		// Whether the requests are done sequentially. The offset will start
		// at 0 and wrap around if needed. Otherwise, requests are random but
		// will be page-aligned.
		sequential bool
	}{
		{
			name:           "SeqWrite_InitiallyEmpty",
			initialDensity: 0,
			readFraction:   0,
			sequential:     true,
		},
		{
			name:           "RandReadWrite_InitiallyEmpty",
			initialDensity: 0,
			readFraction:   0.9,
			sequential:     false,
		},
		{
			name:           "RandReadWrite_InitiallyHalfFull",
			initialDensity: 0.5,
			readFraction:   0.9,
			sequential:     false,
		},
	} {
		order := "Random"
		if test.sequential {
			order = "Sequential"
		}
		name := fmt.Sprintf("%s[D=%.1f,%s:N=%d:R=%.1f/W=%.1f]", test.name, test.initialDensity, order, ioCountPerBenchOp, test.readFraction, 1-test.readFraction)
		b.Run(name, func(b *testing.B) {
			b.StopTimer()
			tmp := testfs.MakeTempDir(b)
			for i := 0; i < b.N; i++ {
				// Set up the initial file and chunk it up into a COW
				f, err := os.CreateTemp(tmp, "")
				require.NoError(b, err)
				err = f.Truncate(diskSize)
				ioBlockSize := ioBlockSize(b, f.Name())

				buf := make([]byte, ioBlockSize*32)
				require.NoError(b, err)
				for off := int64(0); off < diskSize; off += int64(len(buf)) {
					if rand.Float64() >= test.initialDensity {
						continue
					}
					_, err := rand.Read(buf)
					require.NoError(b, err)
					_, err = f.WriteAt(buf, off)
					require.NoError(b, err)
				}
				chunkDir, err := os.MkdirTemp(tmp, "")
				require.NoError(b, err)
				cow, err := copy_on_write.ConvertFileToCOW(ctx, env, f.Name(), chunkSize, chunkDir, "", false)
				require.NoError(b, err)
				err = os.Remove(f.Name())
				require.NoError(b, err)

				// Prepare read/write bufs
				randBuf := make([]byte, ioSize)
				_, err = rand.Read(randBuf)
				require.NoError(b, err)
				readBuf := make([]byte, ioSize)
				off := int64(0)

				b.StartTimer()
				for r := 0; r < ioCountPerBenchOp; r++ {
					if !test.sequential {
						off = rand.Int63n(ioBlockSize)
					}
					if rand.Float64() < test.readFraction {
						_, err := cow.ReadAt(readBuf, off)
						require.NoError(b, err)
					} else {
						_, err := cow.WriteAt(randBuf, off)
						require.NoError(b, err)
					}
					if test.sequential {
						off += ioSize
						off %= diskSize
					}
				}
				b.StopTimer()

				// Clean up so we don't run out of resources mid-bench
				err = cow.Close()
				require.NoError(b, err)
				err = os.RemoveAll(chunkDir)
				require.NoError(b, err)
			}
		})
	}
}

func BenchmarkConvertFileToCOW(b *testing.B) {
	flags.Set(b, "app.log_level", "error")
	log.Configure()

	parentDir := testfs.MakeTempDir(b)
	defer os.RemoveAll(parentDir)
	testfs.MakeDirAll(b, parentDir, b.Name())
	chunkSizeBytes := int64(os.Getpagesize() * 1000) // same as firecracker.cowChunkSizeBytes
	blockSize := ioBlockSize(b, parentDir)
	data := bytes.Repeat([]byte{1, 2, 3}, int(blockSize+117)) // Don't align writes to blocks

	for _, sparsenessRatio := range []float64{0, 0.1, 0.5, 0.9, 0.99, 1} {
		b.Run(fmt.Sprintf("%vSparse", sparsenessRatio), func(b *testing.B) {
			inputPath := filepath.Join(parentDir, b.Name())
			r := rand.New(rand.NewSource(1))
			buf := make([]byte, 100_000_000) // 100MB
			for i := 0; i < len(buf); i += len(data) {
				if r.Float64() >= sparsenessRatio {
					copy(buf[i:], data)
				}
			}
			writeSparseFile(b, inputPath, buf, blockSize)

			b.ResetTimer()
			for range b.N {
				b.StopTimer()
				outputDir, err := os.MkdirTemp(parentDir, "output")
				require.NoError(b, err)
				b.StartTimer()

				_, err = copy_on_write.ConvertFileToCOW(context.Background(), nil, inputPath, chunkSizeBytes, outputDir, "", false)
				require.NoError(b, err)

				b.StopTimer()
				os.RemoveAll(outputDir)
				b.StartTimer()
			}
			b.StopTimer()
			require.NoError(b, os.Remove(inputPath))
		})
	}
}

func testStore(t *testing.T, s interfaces.Store, path string) {
	size, err := s.SizeBytes()
	require.NoError(t, err, "SizeBytes failed")
	require.Equal(t, backingFileSizeBytes, size, "unexpected SizeBytes")

	// Try writing out of bounds; these should all fail.
	for _, bounds := range []struct{ Offset, Length int64 }{
		{Offset: -1, Length: 0},
		{Offset: -1, Length: 1},
		{Offset: -1, Length: 2},
		{Offset: -1, Length: size + 1},
		{Offset: -1, Length: size + 2},
		{Offset: size, Length: 1},
		{Offset: size, Length: 2},
		{Offset: size - 1, Length: 2},
		{Offset: size - 1, Length: 3},
	} {
		msg := fmt.Sprintf("offset=%d length=%d, file_size=%d should fail and return n=0", bounds.Offset, bounds.Length, size)
		b := make([]byte, bounds.Length)
		n, err := s.ReadAt(b, bounds.Offset)
		require.Equal(t, 0, n, "%s", msg)
		require.Error(t, err, "%s", msg)
		n, err = s.WriteAt(b, bounds.Offset)
		require.Equal(t, 0, n, "%s", msg)
		require.Error(t, err, "%s", msg)
	}

	expectedContent := make([]byte, int(size))
	buf := make([]byte, int(size))
	n := 1 + rand.Intn(50)
	for i := 0; i < n; i++ {
		// With equal probability, either (a) read a random range and make sure
		// it matches expectedContent, or (b) write a random range and update
		// our expectedContent for subsequent reads.
		offset, length := randSubslice(len(expectedContent))
		if rand.Float64() > 0.5 {
			n, err := s.ReadAt(buf[:length], int64(offset))
			require.NoError(t, err)
			require.Equal(t, int64(length), int64(n))
			require.Equal(t, expectedContent[offset:offset+length], buf[:length])
		} else {
			_, err := rand.Read(buf[:length])
			require.NoError(t, err)

			n, err := s.WriteAt(buf[:length], int64(offset))
			require.NoError(t, err)
			require.Equal(t, int64(length), int64(n))

			copy(expectedContent[offset:offset+length], buf[:length])
		}
	}
	// Make sure we can sync and close without error.
	err = s.Sync()
	require.NoError(t, err, "Sync failed")

	// Read back the raw backing file using read(2); it should match our
	// expected contents. This is especially important for testing mmap, since
	// Sync() should guarantee that the contents are flushed from memory back
	// to disk.
	if path != "" {
		b, err := os.ReadFile(path)
		require.NoError(t, err)
		require.Equal(t, expectedContent, b)
	}

	err = s.Close()
	require.NoError(t, err, "Close failed")
}

// StoreTester allows testing IO operations on a store.
//
// It keeps an internal buffer that tracks the store's currently expected
// contents. If a Read ever returns different results than our expected buffer,
// the test fails.
//
// It ensures that any IO operations done on the store match what we would see
// in practice. For example, we may see concurrent writes to individual chunks,
// but not concurrent writes to overlapping byte ranges. VBD (for example)
// supports concurrent reads and writes, but the guest will not concurrently
// issue reads and writes for overlapping byte ranges, because unix
// read()/write() operations are not atomic. UFFD faults are handled in a single
// goroutine so those don't support concurrent reads/writes at all.
type StoreTester struct {
	t         *testing.T
	chunkSize int64

	rangeLock *RangeLock
	Store     interfaces.Store

	mu  sync.RWMutex
	buf []byte
}

// NewStoreTester returns a tester for the given store. The store is expected to
// be initially empty (i.e. all zeroes).
func NewStoreTester(t *testing.T, store interfaces.Store) *StoreTester {
	size, err := store.SizeBytes()
	require.NoError(t, err)
	chunkSize := size
	if cow, ok := store.(*copy_on_write.COWStore); ok {
		chunkSize = cow.ChunkSizeBytes()
	}
	return &StoreTester{
		t:         t,
		Store:     store,
		rangeLock: NewRangeLock(t, size),
		buf:       make([]byte, size),
		chunkSize: chunkSize,
	}
}

func (st *StoreTester) randSubslice() (p []byte, off int64) {
	offset, length := randSubslice(len(st.buf))
	// Limit the slice length to 10 chunks to reduce contention and exercise
	// non-overlapping concurrency more. Also, in practice we won't see read() /
	// write() calls spanning several chunks.
	length = min(length, 10*int(st.chunkSize))
	return make([]byte, length), int64(offset)
}

// ReadRandomRange reads a random byte range within the store and fails the test
// if it fails or if the returned bytes do not match the current expected buffer
// contents.
func (st *StoreTester) ReadRandomRange() {
	p, off := st.randSubslice()
	st.read(p, off)
}

// ReadAll reads all bytes and fails the test if the read fails or if the
// returned bytes do not match the current expected buffer contents.
func (st *StoreTester) ReadAll() {
	st.read(make([]byte, len(st.buf)), 0)
}

func (st *StoreTester) read(p []byte, off int64) {
	unlock := st.rangeLock.RLock(off, int64(len(p)))
	defer unlock()
	_, err := st.Store.ReadAt(p, off)
	require.NoError(st.t, err)
	st.mu.RLock()
	defer st.mu.RUnlock()
	if !bytes.Equal(p, st.buf[off:off+int64(len(p))]) {
		require.FailNowf(st.t, "read unexpected bytes", "offset 0x%x, length 0x%x", off, len(p))
	}
}

// WriteRandomRange writes a random byte range within the store and fails the
// test if it fails.
func (st *StoreTester) WriteRandomRange() {
	p, off := st.randSubslice()
	_, err := rand.Read(p)
	require.NoError(st.t, err)
	unlock := st.rangeLock.Lock(off, int64(len(p)))
	defer unlock()
	_, err = st.Store.WriteAt(p, off)
	require.NoError(st.t, err)

	st.mu.Lock()
	defer st.mu.Unlock()
	copy(st.buf[off:], p)
}

// RangeLock implements byte-range locking.
type RangeLock struct {
	t    *testing.T
	path string
}

func NewRangeLock(t *testing.T, size int64) *RangeLock {
	// fcntl flock implements byte range locking when using "open file
	// description" OFD locking, which are a Linux-specific feature. So we just
	// use that for now rather than implement byte-range locking ourselves in
	// Go. (There doesn't appear to be a library available at the time of
	// writing).
	if runtime.GOOS != "linux" {
		t.Skipf("RangeLock is not implemented on %q", runtime.GOOS)
	}
	require.NotEqual(t, 0, cow_cgo_testutil.F_OFD_SETLKW, "sanity check: F_OFD_SETLKW should be defined")
	path := testfs.MakeTempFile(t, testfs.MakeTempDir(t), "")
	err := os.Truncate(path, size)
	require.NoError(t, err)
	return &RangeLock{t: t, path: path}
}

func (rl *RangeLock) RLock(offset, length int64) (unlock func()) {
	return rl.flock(offset, length, syscall.F_RDLCK)
}

func (rl *RangeLock) Lock(offset, length int64) (unlock func()) {
	return rl.flock(offset, length, syscall.F_WRLCK)
}

func (rl *RangeLock) flock(offset, length int64, typ int16) (unlock func()) {
	f, err := os.OpenFile(rl.path, os.O_RDWR, 0)
	require.NoError(rl.t, err)
	const SEEK_SET = 0
	params := &syscall.Flock_t{
		Type:   typ,
		Start:  offset,
		Len:    length,
		Whence: SEEK_SET,
	}
	err = syscall.FcntlFlock(f.Fd(), cow_cgo_testutil.F_OFD_SETLKW, params)
	require.NoError(rl.t, err)
	return func() {
		err := f.Close()
		require.NoError(rl.t, err)
	}
}

func randBytes(t *testing.T, n int) []byte {
	b := make([]byte, n)
	_, err := rand.Read(b)
	require.NoError(t, err)
	return b
}

// Picks a uniform random subslice of a slice with a given length.
// Returns the offset and length of the subslice.
func randSubslice(sliceLength int) (offset, length int) {
	length = rand.Intn(sliceLength + 1)
	offset = rand.Intn(sliceLength - length + 1)
	return
}

func newMmap(t *testing.T) (*copy_on_write.Mmap, string) {
	ctx := context.Background()
	env := testenv.GetTestEnv(t)

	root := testfs.MakeTempDir(t)
	const offset = 0
	path := filepath.Join(root, fmt.Sprintf("%d", offset))
	err := os.WriteFile(path, make([]byte, backingFileSizeBytes), 0644)
	require.NoError(t, err, "write empty file")

	f, err := os.OpenFile(path, os.O_RDWR, 0)
	require.NoError(t, err)
	defer f.Close()
	s, err := f.Stat()
	require.NoError(t, err)

	mmap, err := copy_on_write.NewMmapFd(ctx, env, root, false /*=dirty*/, int(f.Fd()), int(s.Size()), offset, snaputil.ChunkSourceLocalFile, "", false)
	require.NoError(t, err)
	return mmap, path
}

func makeTempFile(t *testing.T, content []byte) string {
	root := testfs.MakeTempDir(t)
	path := filepath.Join(root, "f")
	err := os.WriteFile(path, content, 0644)
	require.NoError(t, err, "write empty file")
	return path
}

func makeEmptyTempFile(t *testing.T, sizeBytes int64) string {
	return makeTempFile(t, make([]byte, sizeBytes))
}

// writeSparseFile writes only the data blocks from b to the given path, so
// that the physical size of the file is minimal while still representing the
// same underlying bytes.
func writeSparseFile(t testing.TB, path string, b []byte, ioBlockSize int64) {
	f, err := os.Create(path)
	require.NoError(t, err)
	defer f.Close()
	// Truncate to ensure the file has the correct size in case it ends with one
	// or more empty blocks.
	err = f.Truncate(int64(len(b)))
	require.NoError(t, err)
	for off := int64(0); off < int64(len(b)); off += ioBlockSize {
		data := b[off:]
		if int64(len(data)) > ioBlockSize {
			data = data[:ioBlockSize]
		}
		if copy_on_write.IsEmptyOrAllZero(data) {
			continue
		}
		_, err := f.WriteAt(data, off)
		require.NoError(t, err)
	}
}

func numIOBlocks(t *testing.T, path string) int64 {
	s := &syscall.Stat_t{}
	err := syscall.Stat(path, s)
	require.NoError(t, err)
	// stat() always uses 512 bytes for its block size, which may not match the
	// IO block size (block size used by the filesystem).
	// See https://askubuntu.com/a/1308745
	statBlocksPerIOBlock := int64(s.Blksize) / 512
	return s.Blocks / statBlocksPerIOBlock
}

func ioBlockSize(t testing.TB, path string) int64 {
	st := &syscall.Stat_t{}
	err := syscall.Stat(path, st)
	require.NoError(t, err)
	return int64(st.Blksize)
}

func concatBytes(chunks ...[]byte) []byte {
	var out []byte
	for _, c := range chunks {
		out = append(out, c...)
	}
	return out
}
