github.com/NVIDIA/aistore@v1.3.23-0.20240517131212-7df6609be51d/fs/walkbck.go (about)

     1  // Package fs provides mountpath and FQN abstractions and methods to resolve/map stored content
     2  /*
     3   * Copyright (c) 2018-2024, NVIDIA CORPORATION. All rights reserved.
     4   */
     5  package fs
     6  
     7  import (
     8  	"container/heap"
     9  	"context"
    10  
    11  	"github.com/NVIDIA/aistore/cmn"
    12  	"github.com/NVIDIA/aistore/cmn/debug"
    13  	"golang.org/x/sync/errgroup"
    14  )
    15  
    16  // provides default low-level `jogger` to traverse a bucket - a poor-man's joggers of sorts
    17  // for those (very few) clients that don't have their own custom implementation
    18  
    19  type WalkBckOpts struct {
    20  	ValidateCb walkFunc // should return filepath.SkipDir to skip directory without an error
    21  	WalkOpts
    22  }
    23  
    24  // internals
    25  type (
    26  	joggerBck struct {
    27  		workCh   chan *wbe
    28  		mi       *Mountpath
    29  		validate walkFunc
    30  		ctx      context.Context
    31  		opts     WalkOpts
    32  	}
    33  	wbe struct { // walk bck entry
    34  		dirEntry DirEntry
    35  		fqn      string
    36  	}
    37  	wbeInfo struct {
    38  		dirEntry DirEntry
    39  		fqn      string
    40  		objName  string
    41  		mpathIdx int
    42  	}
    43  	wbeHeap []wbeInfo
    44  )
    45  
    46  func WalkBck(opts *WalkBckOpts) error {
    47  	debug.Assert(opts.Mi == nil && opts.Sorted) // TODO: support `opts.Sorted == false`
    48  	var (
    49  		avail      = GetAvail()
    50  		l          = len(avail)
    51  		joggers    = make([]*joggerBck, l)
    52  		group, ctx = errgroup.WithContext(context.Background())
    53  		idx        int
    54  	)
    55  	for _, mi := range avail {
    56  		workCh := make(chan *wbe, mpathQueueSize)
    57  		jg := &joggerBck{
    58  			workCh:   workCh,
    59  			mi:       mi,
    60  			validate: opts.ValidateCb,
    61  			ctx:      ctx,
    62  			opts:     opts.WalkOpts,
    63  		}
    64  		jg.opts.Callback = jg.cb // --> jg.validate --> opts.ValidateCb
    65  		jg.opts.Mi = mi
    66  		joggers[idx] = jg
    67  		idx++
    68  	}
    69  
    70  	for i := range l {
    71  		group.Go(joggers[i].walk)
    72  	}
    73  	group.Go(func() error {
    74  		h := &wbeHeap{}
    75  		heap.Init(h)
    76  
    77  		for i := range l {
    78  			if wbe, ok := <-joggers[i].workCh; ok {
    79  				heap.Push(h, wbeInfo{mpathIdx: i, fqn: wbe.fqn, dirEntry: wbe.dirEntry})
    80  			}
    81  		}
    82  		for h.Len() > 0 {
    83  			v := heap.Pop(h)
    84  			info := v.(wbeInfo)
    85  			if err := opts.Callback(info.fqn, info.dirEntry); err != nil {
    86  				return err
    87  			}
    88  			if wbe, ok := <-joggers[info.mpathIdx].workCh; ok {
    89  				heap.Push(h, wbeInfo{mpathIdx: info.mpathIdx, fqn: wbe.fqn, dirEntry: wbe.dirEntry})
    90  			}
    91  		}
    92  		return nil
    93  	})
    94  
    95  	return group.Wait()
    96  }
    97  
    98  ///////////////
    99  // joggerBck //
   100  ///////////////
   101  
   102  func (jg *joggerBck) walk() (err error) {
   103  	err = Walk(&jg.opts)
   104  	close(jg.workCh)
   105  	return
   106  }
   107  
   108  func (jg *joggerBck) cb(fqn string, de DirEntry) error {
   109  	const tag = "fs-walk-bck-mpath"
   110  	select {
   111  	case <-jg.ctx.Done():
   112  		return cmn.NewErrAborted(jg.mi.String(), tag, nil)
   113  	default:
   114  		break
   115  	}
   116  	if jg.validate != nil {
   117  		if err := jg.validate(fqn, de); err != nil {
   118  			// If err != filepath.SkipDir, the Walk will propagate the error to group.Go.
   119  			// Context will be canceled, which then will terminate all running goroutines.
   120  			return err
   121  		}
   122  	}
   123  	if de.IsDir() {
   124  		return nil
   125  	}
   126  	select {
   127  	case <-jg.ctx.Done():
   128  		return cmn.NewErrAborted(jg.mi.String(), tag, nil)
   129  	case jg.workCh <- &wbe{de, fqn}:
   130  		return nil
   131  	}
   132  }
   133  
   134  /////////////
   135  // wbeHeap //
   136  /////////////
   137  
   138  func (h wbeHeap) Len() int           { return len(h) }
   139  func (h wbeHeap) Less(i, j int) bool { return h[i].objName < h[j].objName }
   140  func (h wbeHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
   141  
   142  func (h *wbeHeap) Push(x any) {
   143  	var (
   144  		parsed ParsedFQN
   145  		info   = x.(wbeInfo)
   146  	)
   147  	debug.Assert(info.objName == "")
   148  	if err := parsed.Init(info.fqn); err != nil {
   149  		return
   150  	}
   151  	info.objName = parsed.ObjName
   152  	*h = append(*h, info)
   153  }
   154  
   155  func (h *wbeHeap) Pop() any {
   156  	old := *h
   157  	n := len(old)
   158  	x := old[n-1]
   159  	*h = old[0 : n-1]
   160  	return x
   161  }