feat: new UnixFS{File,Directory} with options pattern

flag future deprecation of Generate*

Author: rvagg

testutil/generator.go

@@ -19,31 +19,225 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-// GenerateFile generates a random unixfs file of the given size, storing the
-// blocks in the provided LinkSystem and returns a DirEntry representation of
+type options struct {
+	randReader     io.Reader
+	shardBitwidth  int
+	chunker        string
+	dirname        string
+	childGenerator ChildGeneratorFn
+	shardThisDir   bool // a private option used internally to randomly switch on sharding at this current level
+}
+
+// Option is a functional option for the Generate* functions.
+type Option func(*options)
+
+// WithRandReader sets the random reader used by the Generate* functions.
+func WithRandReader(randReader io.Reader) Option {
+	return func(o *options) {
+		o.randReader = randReader
+	}
+}
+
+// WithShardBitwidth sets the shard fanout for the generated directory. By
+// default directories are not sharded. Set to 8 to use the default sharding
+// fanout value of 256. Set to a lower value, such as 4, to increase the
+// probability of collisions and therefore greater depth for smaller number of
+// files.
+func WithShardBitwidth(bitwidth int) Option {
+	return func(o *options) {
+		o.shardBitwidth = bitwidth
+	}
+}
+
+// WithChunker sets the chunker used by the Generate* functions. By default
+// files are chunked using the "size-256144" chunker.
+// The "size-256144" chunker will result in splitting bytes at 256144b
+// boundaries. See https://pkg.go.dev/github.com/ipfs/go-ipfs-chunker#FromString
+// for more information on options available.
+func WithChunker(chunker string) Option {
+	return func(o *options) {
+		o.chunker = chunker
+	}
+}
+
+// WithDirname sets the directory name used by UnixFSDirectory where a root
+// directory name is required.
+func WithDirname(dirname string) Option {
+	return func(o *options) {
+		o.dirname = dirname
+	}
+}
+
+// ChildGeneratorFn is a function that generates a child DirEntry for a
+// directory. It is used by UnixFSDirectory where control over the direct
+// children of a directory is required. Return nil to stop generating children.
+type ChildGeneratorFn func(name string) (*DirEntry, error)
+
+// WithChildGenerator sets the child generator used by UnixFSDirectory control
+// over the direct children of a directory is required.
+func WithChildGenerator(childGenerator ChildGeneratorFn) Option {
+	return func(o *options) {
+		o.childGenerator = childGenerator
+	}
+}
+
+// shardThisDir is a private internal option
+func shardThisDir(b bool) Option {
+	return func(o *options) {
+		o.shardThisDir = b
+	}
+}
+
+func applyOptions(opts []Option) *options {
+	o := &options{
+		randReader:    rand.Reader,
+		shardBitwidth: 0,
+		chunker:       "size-256144",
+		shardThisDir:  true,
+	}
+	for _, opt := range opts {
+		opt(o)
+	}
+	return o
+}
+
+// Generate a file of `size` random bytes, packaged into UnixFS structure,
+// stored in the provided LinkSystem and returns a DirEntry representation of
 // the file.
-func GenerateFile(t require.TestingT, linkSys *linking.LinkSystem, randReader io.Reader, size int) DirEntry {
-	// a file of `size` random bytes, packaged into unixfs DAGs, stored in the remote blockstore
-	delimited := io.LimitReader(randReader, int64(size))
+func UnixFSFile(lsys linking.LinkSystem, size int, opt ...Option) (DirEntry, error) {
+	o := applyOptions(opt)
+	delimited := io.LimitReader(o.randReader, int64(size))
 	var buf bytes.Buffer
 	buf.Grow(size)
 	delimited = io.TeeReader(delimited, &buf)
-	// "size-256144" sets the chunker, splitting bytes at 256144b boundaries
 	cids := make([]cid.Cid, 0)
 	var undo func()
-	linkSys.StorageWriteOpener, undo = cidCollector(linkSys, &cids)
+	lsys.StorageWriteOpener, undo = cidCollector(lsys, &cids)
 	defer undo()
-	root, gotSize, err := builder.BuildUnixFSFile(delimited, "size-256144", linkSys)
-	require.NoError(t, err)
-	srcData := buf.Bytes()
-	rootCid := root.(cidlink.Link).Cid
+	root, gotSize, err := builder.BuildUnixFSFile(delimited, o.chunker, &lsys)
+	if err != nil {
+		return DirEntry{}, err
+	}
 	return DirEntry{
 		Path:     "",
-		Content:  srcData,
-		Root:     rootCid,
+		Content:  buf.Bytes(),
+		Root:     root.(cidlink.Link).Cid,
 		SelfCids: cids,
 		TSize:    uint64(gotSize),
+	}, nil
+}
+
+// UnixFSDirectory generates a random UnixFS directory that aims for the
+// requested targetSize (in bytes, although it is likely to fall somewhere
+// under this number), storing the blocks in the provided LinkSystem and
+// returns a DirEntry representation of the directory.
+//
+// If the WithDirname option is not set, the directory will be built as a root
+// directory. If the dirname option is set, the directory will be built as a
+// child directory.
+//
+// If the WithChildGenerator option is not set, the targetSize will be
+// ignored and all sizing control will be delegated to the child generator.
+func UnixFSDirectory(lsys linking.LinkSystem, targetSize int, opts ...Option) (DirEntry, error) {
+	o := applyOptions(opts)
+
+	var curSize int
+	targetFileSize := targetSize / 16
+
+	children := make([]DirEntry, 0)
+
+	childGenerator := func(name string) (*DirEntry, error) {
+		for curSize < targetSize {
+			switch rndInt(o.randReader, 6) {
+			case 0: // 1 in 6 chance of finishing this directory if not at root
+				if o.dirname != "" && len(children) > 0 {
+					curSize = targetSize // not really, but we're done with this directory
+				} // else at the root we don't get to finish early
+			case 1: // 1 in 6 chance of making a new directory
+				if targetSize-curSize <= 1024 { // don't make tiny directories
+					continue
+				}
+				so := append(opts, WithDirname(o.dirname+"/"+name), shardThisDir(rndInt(o.randReader, 6) == 0))
+				child, err := UnixFSDirectory(lsys, targetSize-curSize, so...)
+				if err != nil {
+					return nil, err
+				}
+				children = append(children, child)
+				curSize += int(child.TSize)
+			default: // 4 in 6 chance of making a new file
+				var size int
+				for size == 0 { // don't make empty files
+					sizeB, err := rand.Int(o.randReader, big.NewInt(int64(targetFileSize)))
+					if err != nil {
+						return nil, err
+					}
+					size = int(sizeB.Int64())
+					if size > targetSize-curSize {
+						size = targetSize - curSize
+					}
+				}
+				entry, err := UnixFSFile(lsys, size, opts...)
+				if err != nil {
+					return nil, err
+				}
+				var name string
+				entry.Path = o.dirname + "/" + name
+				curSize += size
+				return &entry, nil
+			}
+		}
+		return nil, nil
+	}
+
+	if o.childGenerator != nil {
+		childGenerator = o.childGenerator
+	}
+
+	for {
+		var name string
+		for {
+			var err error
+			name, err = namegen.RandomDirectoryName(o.randReader)
+			if err != nil {
+				return DirEntry{}, err
+			}
+			if !isDupe(children, name) {
+				break
+			}
+		}
+
+		child, err := childGenerator(o.dirname + "/" + name)
+		if err != nil {
+			return DirEntry{}, err
+		}
+		if child == nil {
+			break
+		}
+		children = append(children, *child)
+	}
+
+	var shardBitwidth int
+	if o.shardThisDir {
+		shardBitwidth = o.shardBitwidth
+	}
+	dirEntry, err := packDirectory(lsys, children, shardBitwidth)
+	if err != nil {
+		return DirEntry{}, err
 	}
+	dirEntry.Path = o.dirname
+	return dirEntry, nil
+}
+
+// GenerateFile generates a random unixfs file of the given size, storing the
+// blocks in the provided LinkSystem and returns a DirEntry representation of
+// the file.
+//
+// This function will be deprecated in a future release, use UnixFSFile()
+// instead.
+func GenerateFile(t require.TestingT, linkSys *linking.LinkSystem, randReader io.Reader, size int) DirEntry {
+	dirEntry, err := UnixFSFile(*linkSys, size, WithRandReader(randReader))
+	require.NoError(t, err)
+	return dirEntry
 }
 
 // GenerateDirectory generates a random UnixFS directory that aims for the
@@ -53,13 +247,19 @@ func GenerateFile(t require.TestingT, linkSys *linking.LinkSystem, randReader io
 // the root directory will be built as HAMT sharded (with a low "width" to
 // maximise the chance of collisions and therefore greater depth for smaller
 // number of files).
+//
+// This function will be deprecated in a future release, use UnixFSDirectory()
+// instead.
 func GenerateDirectory(t require.TestingT, linkSys *linking.LinkSystem, randReader io.Reader, targetSize int, rootSharded bool) DirEntry {
 	return GenerateDirectoryFrom(t, linkSys, randReader, targetSize, "", rootSharded)
 }
 
 // GenerateDirectoryFrom is the same as GenerateDirectory but allows the caller
 // to specify a directory path to start from. This is useful for generating
 // nested directories.
+//
+// This function will be deprecated in a future release, use UnixFSDirectory()
+// instead.
 func GenerateDirectoryFrom(
 	t require.TestingT,
 	linkSys *linking.LinkSystem,
@@ -124,46 +324,26 @@ func GenerateDirectoryFrom(
 	return dirEntry
 }
 
-// GenerateDirectoryFor will build a directory with name `dir` whose children
-// are generated by the caller using the provided mkchild function. This is
-// useful for generating directories with a specific structure where you need
-// to have greater control over the children.
-func GenerateDirectoryFor(
-	t require.TestingT,
-	linkSys *linking.LinkSystem,
-	randReader io.Reader,
-	dir string,
-	sharded bool,
-	mkchild func(name string) *DirEntry,
-) DirEntry {
-	children := make([]DirEntry, 0)
-	for {
-		var name string
-		for {
-			var err error
-			name, err = namegen.RandomDirectoryName(randReader)
-			require.NoError(t, err)
-			if !isDupe(children, name) {
-				break
-			}
-		}
-		child := mkchild(dir + "/" + name)
-		if child == nil {
-			break
-		}
-		children = append(children, *child)
-	}
-	dirEntry := BuildDirectory(t, linkSys, children, sharded)
-	dirEntry.Path = dir
-	return dirEntry
-}
-
 // BuildDirectory builds a directory from the given children, storing the
 // blocks in the provided LinkSystem and returns a DirEntry representation of
 // the directory. If sharded is true, the root directory will be built as HAMT
 // sharded (with a low "width" to maximise the chance of collisions and
 // therefore greater depth for smaller number of files).
+//
+// This function will be deprecated in a future release. Currently there is no
+// direct replacement.
 func BuildDirectory(t require.TestingT, linkSys *linking.LinkSystem, children []DirEntry, sharded bool) DirEntry {
+	var bitWidth int
+	if sharded {
+		// fanout of 16, quite small to increase collision probability so we actually get sharding
+		bitWidth = 4
+	}
+	dirEnt, err := packDirectory(*linkSys, children, bitWidth)
+	require.NoError(t, err)
+	return dirEnt
+}
+
+func packDirectory(lsys linking.LinkSystem, children []DirEntry, bitWidth int) (DirEntry, error) {
 	// create stable sorted children, which should match the encoded form
 	// in dag-pb
 	sort.Slice(children, func(i, j int) bool {
@@ -175,25 +355,31 @@ func BuildDirectory(t require.TestingT, linkSys *linking.LinkSystem, children []
 		paths := strings.Split(child.Path, "/")
 		name := paths[len(paths)-1]
 		lnk, err := builder.BuildUnixFSDirectoryEntry(name, int64(child.TSize), cidlink.Link{Cid: child.Root})
-		require.NoError(t, err)
+		if err != nil {
+			return DirEntry{}, err
+		}
 		dirLinks = append(dirLinks, lnk)
 	}
 	cids := make([]cid.Cid, 0)
 	var undo func()
-	linkSys.StorageWriteOpener, undo = cidCollector(linkSys, &cids)
+	lsys.StorageWriteOpener, undo = cidCollector(lsys, &cids)
 	defer undo()
 	var root ipld.Link
 	var size uint64
 	var err error
-	if sharded {
-		// node arity of 16, quite small to increase collision probability so we actually get sharding
-		const width = 16
+	if bitWidth > 0 {
+		// width is 2^bitWidth
+		width := 2 << bitWidth
 		const hasher = multihash.MURMUR3X64_64
-		root, size, err = builder.BuildUnixFSShardedDirectory(width, hasher, dirLinks, linkSys)
-		require.NoError(t, err)
+		root, size, err = builder.BuildUnixFSShardedDirectory(width, hasher, dirLinks, &lsys)
+		if err != nil {
+			return DirEntry{}, err
+		}
 	} else {
-		root, size, err = builder.BuildUnixFSDirectory(dirLinks, linkSys)
-		require.NoError(t, err)
+		root, size, err = builder.BuildUnixFSDirectory(dirLinks, &lsys)
+		if err != nil {
+			return DirEntry{}, err
+		}
 	}
 
 	return DirEntry{
@@ -202,7 +388,7 @@ func BuildDirectory(t require.TestingT, linkSys *linking.LinkSystem, children []
 		SelfCids: cids,
 		TSize:    size,
 		Children: children,
-	}
+	}, nil
 }
 
 func rndInt(randReader io.Reader, max int) int {
@@ -213,7 +399,7 @@ func rndInt(randReader io.Reader, max int) int {
 	return int(coin.Int64())
 }
 
-func cidCollector(ls *ipld.LinkSystem, cids *[]cid.Cid) (ipld.BlockWriteOpener, func()) {
+func cidCollector(ls ipld.LinkSystem, cids *[]cid.Cid) (ipld.BlockWriteOpener, func()) {
 	swo := ls.StorageWriteOpener
 	return func(linkCtx ipld.LinkContext) (io.Writer, ipld.BlockWriteCommitter, error) {
 			w, c, err := swo(linkCtx)