ipld/unixfs/io: add WithMaxLinks() to directories

The new option sets effectively a maximum width to the directories.

Currently, on a dynamic directory, the switching to-from Basic to HAMT is
controlled by HAMTShardingSize, and the width of the HAMT shards by DefaultShardWidth.

When WithMaxLinks() is specified, the switching is additionally controlled by
the number of links exceeding the limit. In that case, MaxLinks is used as ShardWidth.

The directory can only be converted back to BasicDirectory when the total
number of links is below the limit.

Backwards compatibility is kept and tests have been added.

Note that when setting MaxLinks to a high number, it is possible that we still
suffer automatic conversion to HAMT before hitting MaxLinks, if the estimated
directory size is above 256KiB (as before).

Author: hsanjuan

ipld/unixfs/io/directory.go

@@ -69,6 +69,39 @@ type Directory interface {
 
 	// GetCidBuilder returns the CID Builder used.
 	GetCidBuilder() cid.Builder
+
+	// SetMaxLinks sets the max width of the Directory.
+	SetMaxLinks(n int)
+}
+
+type DirectoryOption func(Directory)
+
+// WithMaxLinks stablishes the max number of links allowed for a directory:
+//
+// - On Dynamic directories using a BasicDirectory, it can trigger conversion
+// to HAMT when set and exceeded. The subsequent HAMT nodes will use MaxLinks
+// as ShardWidth. Conversion can happen too based on HAMTShardingSize.
+//
+// - On Dynamic directories using a HAMTDirectory, it can trigger conversion
+// to BasicDirectory when the number of links is below MaxLinks (and
+// HAMTShardingSize allows).
+//
+// - On Basic directories, it causes an error when adding more than MaxLinks
+// children.
+//
+// - On HAMT directories, it sets the ShardWidth, otherwise DefaultShardWidth
+// is used.
+func WithMaxLinks(n int) DirectoryOption {
+	return func(d Directory) {
+		d.SetMaxLinks(n)
+	}
+}
+
+// WithCidBuilder sets the CidBuilder for new Directories.
+func WithCidBuilder(cb cid.Builder) DirectoryOption {
+	return func(d Directory) {
+		d.SetCidBuilder(cb)
+	}
 }
 
 // TODO: Evaluate removing `dserv` from this layer and providing it in MFS.
@@ -96,6 +129,11 @@ type BasicDirectory struct {
 	// (We maintain this value up to date even if the HAMTShardingSize is off
 	// since potentially the option could be activated on the fly.)
 	estimatedSize int
+	totalLinks    int
+
+	// opts
+	// maxNumberOfLinks. If set, can trigger conversion to HAMT directory.
+	maxLinks int
 }
 
 // HAMTDirectory is the HAMT implementation of `Directory`.
@@ -104,30 +142,82 @@ type HAMTDirectory struct {
 	shard *hamt.Shard
 	dserv ipld.DAGService
 
+	// opts
+	maxLinks   int
+	cidBuilder cid.Builder
+
 	// Track the changes in size by the AddChild and RemoveChild calls
 	// for the HAMTShardingSize option.
 	sizeChange int
+	totalLinks int
 }
 
-func newEmptyBasicDirectory(dserv ipld.DAGService) *BasicDirectory {
-	return newBasicDirectoryFromNode(dserv, format.EmptyDirNode())
+func newEmptyBasicDirectory(dserv ipld.DAGService, opts ...DirectoryOption) *BasicDirectory {
+	return newBasicDirectoryFromNode(dserv, format.EmptyDirNode(), opts...)
 }
 
-func newBasicDirectoryFromNode(dserv ipld.DAGService, node *mdag.ProtoNode) *BasicDirectory {
+func newBasicDirectoryFromNode(dserv ipld.DAGService, node *mdag.ProtoNode, opts ...DirectoryOption) *BasicDirectory {
 	basicDir := new(BasicDirectory)
 	basicDir.node = node
 	basicDir.dserv = dserv
 
+	for _, o := range opts {
+		o(basicDir)
+	}
+
 	// Scan node links (if any) to restore estimated size.
-	basicDir.computeEstimatedSize()
+	basicDir.computeEstimatedSizeAndTotalLinks()
 
 	return basicDir
 }
 
-// NewDirectory returns a Directory implemented by DynamicDirectory
-// containing a BasicDirectory that can be converted to a HAMTDirectory.
-func NewDirectory(dserv ipld.DAGService) Directory {
-	return &DynamicDirectory{newEmptyBasicDirectory(dserv)}
+func newEmptyHAMTDirectory(dserv ipld.DAGService, sizeChange int, opts ...DirectoryOption) (*HAMTDirectory, error) {
+	dir := new(HAMTDirectory)
+	dir.dserv = dserv
+	dir.sizeChange = sizeChange
+	dir.maxLinks = 0
+
+	for _, opt := range opts {
+		opt(dir)
+	}
+
+	// If user did not set maxLinks, we leave it unset, in case this
+	// folder is converted to basic folder at some point.
+	maxLinks := DefaultShardWidth
+	if dir.maxLinks > 0 {
+		maxLinks = dir.maxLinks
+	}
+
+	shard, err := hamt.NewShard(dir.dserv, maxLinks)
+	if err != nil {
+		return nil, err
+	}
+	shard.SetCidBuilder(dir.cidBuilder)
+	dir.shard = shard
+
+	return dir, nil
+}
+
+func newHAMTDirectoryFromNode(dserv ipld.DAGService, node ipld.Node) (*HAMTDirectory, error) {
+	dir := new(HAMTDirectory)
+	dir.dserv = dserv
+	dir.sizeChange = 0
+
+	shard, err := hamt.NewHamtFromDag(dserv, node)
+	if err != nil {
+		return nil, err
+	}
+	dir.shard = shard
+	dir.totalLinks = len(node.Links())
+
+	return dir, nil
+}
+
+// NewDirectory returns a Directory implemented by DynamicDirectory containing
+// a BasicDirectory that automatically converts to a from a HAMTDirectory
+// based on HAMTShardingSize and MaxLinks (when set).
+func NewDirectory(dserv ipld.DAGService, opts ...DirectoryOption) Directory {
+	return &DynamicDirectory{newEmptyBasicDirectory(dserv, opts...)}
 }
 
 // ErrNotADir implies that the given node was not a unixfs directory
@@ -150,21 +240,26 @@ func NewDirectoryFromNode(dserv ipld.DAGService, node ipld.Node) (Directory, err
 	case format.TDirectory:
 		return &DynamicDirectory{newBasicDirectoryFromNode(dserv, protoBufNode.Copy().(*mdag.ProtoNode))}, nil
 	case format.THAMTShard:
-		shard, err := hamt.NewHamtFromDag(dserv, node)
+		hamtDir, err := newHAMTDirectoryFromNode(dserv, node)
 		if err != nil {
 			return nil, err
 		}
-		return &DynamicDirectory{&HAMTDirectory{shard, dserv, 0}}, nil
+		return &DynamicDirectory{hamtDir}, nil
 	}
 
 	return nil, ErrNotADir
 }
 
-func (d *BasicDirectory) computeEstimatedSize() {
+func (d *BasicDirectory) SetMaxLinks(n int) {
+	d.maxLinks = n
+}
+
+func (d *BasicDirectory) computeEstimatedSizeAndTotalLinks() {
 	d.estimatedSize = 0
 	// err is just breaking the iteration and we always return nil
 	_ = d.ForEachLink(context.TODO(), func(l *ipld.Link) error {
 		d.addToEstimatedSize(l.Name, l.Cid)
+		d.totalLinks++
 		return nil
 	})
 	// ForEachLink will never fail traversing the BasicDirectory
@@ -181,7 +276,7 @@ func (d *BasicDirectory) removeFromEstimatedSize(name string, linkCid cid.Cid) {
 		// Something has gone very wrong. Log an error and recompute the
 		// size from scratch.
 		log.Error("BasicDirectory's estimatedSize went below 0")
-		d.computeEstimatedSize()
+		d.computeEstimatedSizeAndTotalLinks()
 	}
 }
 
@@ -219,7 +314,15 @@ func (d *BasicDirectory) needsToSwitchToHAMTDir(name string, nodeToAdd ipld.Node
 		operationSizeChange += linksize.LinkSizeFunction(name, nodeToAdd.Cid())
 	}
 
-	return d.estimatedSize+operationSizeChange >= HAMTShardingSize, nil
+	switchShardingSize := d.estimatedSize+operationSizeChange >= HAMTShardingSize
+	switchMaxLinks := false
+	// We should switch if we have reached maxLinks and a new link is being
+	// added and maxLinks is valid for shardWidth.
+	if nodeToAdd != nil && entryToRemove == nil && validShardWidth(d.maxLinks) &&
+		(d.totalLinks+1) > d.maxLinks {
+		switchMaxLinks = true
+	}
+	return switchShardingSize || switchMaxLinks, nil
 }
 
 // addLinkChild adds the link as an entry to this directory under the given
@@ -231,12 +334,20 @@ func (d *BasicDirectory) addLinkChild(ctx context.Context, name string, link *ip
 	if err != nil && err != os.ErrNotExist {
 		return err
 	}
+	if err == nil { // existed
+		d.totalLinks--
+	}
+
+	if d.maxLinks > 0 && d.totalLinks+1 > d.maxLinks {
+		return errors.New("BasicDirectory: cannot add child: maxLinks reached")
+	}
 
 	err = d.node.AddRawLink(name, link)
 	if err != nil {
 		return err
 	}
 	d.addToEstimatedSize(name, link.Cid)
+	d.totalLinks++
 	return nil
 }
 
@@ -304,6 +415,7 @@ func (d *BasicDirectory) RemoveChild(ctx context.Context, name string) error {
 
 	// The name actually existed so we should update the estimated size.
 	d.removeFromEstimatedSize(link.Name, link.Cid)
+	d.totalLinks--
 
 	return d.node.RemoveNodeLink(name)
 	// GetNodeLink didn't return ErrLinkNotFound so this won't fail with that
@@ -321,30 +433,38 @@ func (d *BasicDirectory) GetCidBuilder() cid.Builder {
 }
 
 // switchToSharding returns a HAMT implementation of this directory.
-func (d *BasicDirectory) switchToSharding(ctx context.Context) (*HAMTDirectory, error) {
-	hamtDir := new(HAMTDirectory)
-	hamtDir.dserv = d.dserv
-
-	shard, err := hamt.NewShard(d.dserv, DefaultShardWidth)
+func (d *BasicDirectory) switchToSharding(ctx context.Context, opts ...DirectoryOption) (*HAMTDirectory, error) {
+	hamtDir, err := newEmptyHAMTDirectory(d.dserv, 0, opts...)
 	if err != nil {
 		return nil, err
 	}
-	shard.SetCidBuilder(d.node.CidBuilder())
-	hamtDir.shard = shard
 
 	for _, lnk := range d.node.Links() {
 		err = hamtDir.shard.SetLink(ctx, lnk.Name, lnk)
 		if err != nil {
 			return nil, err
 		}
+		hamtDir.totalLinks++
 	}
 
 	return hamtDir, nil
 }
 
 // SetCidBuilder implements the `Directory` interface.
 func (d *HAMTDirectory) SetCidBuilder(builder cid.Builder) {
-	d.shard.SetCidBuilder(builder)
+	d.cidBuilder = builder
+	if d.shard != nil {
+		d.shard.SetCidBuilder(builder)
+	}
+}
+
+// SetMaxLinks set the max number of links in a HAMTDirectory. It overrides
+// DefaultShardWidth when present. Needs to be a power of two (shard entry
+// size) and multiple of 8 (bitfield size).
+func (d *HAMTDirectory) SetMaxLinks(maxLinks int) {
+	if validShardWidth(maxLinks) {
+		d.maxLinks = maxLinks
+	}
 }
 
 // AddChild implements the `Directory` interface.
@@ -358,6 +478,9 @@ func (d *HAMTDirectory) AddChild(ctx context.Context, name string, nd ipld.Node)
 		d.removeFromSizeChange(oldChild.Name, oldChild.Cid)
 	}
 	d.addToSizeChange(name, nd.Cid())
+	if oldChild == nil {
+		d.totalLinks++
+	}
 	return nil
 }
 
@@ -396,6 +519,7 @@ func (d *HAMTDirectory) RemoveChild(ctx context.Context, name string) error {
 
 	if oldChild != nil {
 		d.removeFromSizeChange(oldChild.Name, oldChild.Cid)
+		d.totalLinks--
 	}
 
 	return nil
@@ -412,9 +536,10 @@ func (d *HAMTDirectory) GetCidBuilder() cid.Builder {
 }
 
 // switchToBasic returns a BasicDirectory implementation of this directory.
-func (d *HAMTDirectory) switchToBasic(ctx context.Context) (*BasicDirectory, error) {
-	basicDir := newEmptyBasicDirectory(d.dserv)
-	basicDir.SetCidBuilder(d.GetCidBuilder())
+func (d *HAMTDirectory) switchToBasic(ctx context.Context, opts ...DirectoryOption) (*BasicDirectory, error) {
+	// needsToSwichToBasicDir checks d.maxLinks is appropiate. No check is
+	// performed here.
+	basicDir := newEmptyBasicDirectory(d.dserv, opts...)
 
 	err := d.ForEachLink(ctx, func(lnk *ipld.Link) error {
 		err := basicDir.addLinkChild(ctx, lnk.Name, lnk)
@@ -472,14 +597,32 @@ func (d *HAMTDirectory) needsToSwitchToBasicDir(ctx context.Context, name string
 		operationSizeChange += linksize.LinkSizeFunction(name, nodeToAdd.Cid())
 	}
 
-	if d.sizeChange+operationSizeChange >= 0 {
-		// We won't have reduced the HAMT net size.
-		return false, nil
+	// We must switch if size and maxlinks are below threshold
+	canSwitchSize := false
+	// Directory size reduced, perhaps below limit.
+	if d.sizeChange+operationSizeChange < 0 {
+		canSwitchSize, err = d.sizeBelowThreshold(ctx, operationSizeChange)
+		if err != nil {
+			return false, err
+		}
+	}
+
+	canSwitchMaxLinks := true
+	if d.maxLinks > 0 {
+		total := d.totalLinks
+		if nodeToAdd != nil && entryToRemove == nil {
+			total++
+		} else if nodeToAdd == nil && entryToRemove != nil {
+			total--
+		}
+		if total > d.maxLinks {
+			// prevent switching as we would end up with too many links
+			canSwitchMaxLinks = false
+		}
 	}
 
-	// We have reduced the directory size, check if went below the
-	// HAMTShardingSize threshold to trigger a switch.
-	return d.sizeBelowThreshold(ctx, operationSizeChange)
+	return canSwitchSize && canSwitchMaxLinks, nil
+
 }
 
 // Evaluate directory size and a future sizeChange and check if it will be below
@@ -554,7 +697,7 @@ func (d *DynamicDirectory) AddChild(ctx context.Context, name string, nd ipld.No
 		}
 
 		if switchToBasic {
-			basicDir, err := hamtDir.switchToBasic(ctx)
+			basicDir, err := hamtDir.switchToBasic(ctx, WithMaxLinks(hamtDir.maxLinks), WithCidBuilder(hamtDir.GetCidBuilder()))
 			if err != nil {
 				return err
 			}
@@ -578,7 +721,14 @@ func (d *DynamicDirectory) AddChild(ctx context.Context, name string, nd ipld.No
 	if !switchToHAMT {
 		return basicDir.AddChild(ctx, name, nd)
 	}
-	hamtDir, err = basicDir.switchToSharding(ctx)
+
+	maxLinks := DefaultShardWidth
+	// Verify that our maxLinks is usuable for ShardWidth (power of 2, multiple of 8)
+	if validShardWidth(basicDir.maxLinks) {
+		maxLinks = basicDir.maxLinks
+	}
+
+	hamtDir, err = basicDir.switchToSharding(ctx, WithMaxLinks(maxLinks), WithCidBuilder(basicDir.GetCidBuilder()))
 	if err != nil {
 		return err
 	}
@@ -608,14 +758,23 @@ func (d *DynamicDirectory) RemoveChild(ctx context.Context, name string) error {
 		return hamtDir.RemoveChild(ctx, name)
 	}
 
-	basicDir, err := hamtDir.switchToBasic(ctx)
+	// We have not removed the element that violates MaxLinks, so we have to +1 the limit. We -1 below.
+	basicDir, err := hamtDir.switchToBasic(ctx, WithMaxLinks(hamtDir.maxLinks+1), WithCidBuilder(hamtDir.GetCidBuilder()))
 	if err != nil {
 		return err
 	}
 	err = basicDir.RemoveChild(ctx, name)
 	if err != nil {
 		return err
 	}
+
+	basicDir.SetMaxLinks(hamtDir.maxLinks - 1)
 	d.Directory = basicDir
 	return nil
 }
+
+// validShardWidth verifies that the given number is positive, a power of 2
+// and a multiple of 8.
+func validShardWidth(n int) bool {
+	return n > 0 && (n&(n-1)) == 0 && n%8 == 0
+}