github.com/filecoin-project/lassie@v0.23.0/pkg/internal/candidatebuffer/candidatebuffer.go (about)

     1  package candidatebuffer
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  	"time"
     7  
     8  	"github.com/benbjohnson/clock"
     9  	"github.com/filecoin-project/lassie/pkg/types"
    10  )
    11  
    12  type CandidateBuffer struct {
    13  	clock             clock.Clock
    14  	timerCancel       context.CancelFunc
    15  	currentCandidates []types.RetrievalCandidate
    16  	lk                sync.Mutex
    17  	onCandidates      func([]types.RetrievalCandidate)
    18  	afterEach         chan<- struct{}
    19  }
    20  
    21  func NewCandidateBuffer(onCandidates func([]types.RetrievalCandidate), clock clock.Clock) *CandidateBuffer {
    22  	return NewCandidateBufferWithSync(onCandidates, clock, nil)
    23  }
    24  
    25  func NewCandidateBufferWithSync(onCandidates func([]types.RetrievalCandidate), clock clock.Clock, afterEach chan<- struct{}) *CandidateBuffer {
    26  	return &CandidateBuffer{
    27  		onCandidates: onCandidates,
    28  		clock:        clock,
    29  		afterEach:    afterEach,
    30  	}
    31  }
    32  
    33  func (c *CandidateBuffer) clear() []types.RetrievalCandidate {
    34  	c.lk.Lock()
    35  	if c.timerCancel != nil {
    36  		c.timerCancel()
    37  	}
    38  	c.timerCancel = nil
    39  	prevCandidates := c.currentCandidates
    40  	c.currentCandidates = nil
    41  	c.lk.Unlock()
    42  	return prevCandidates
    43  }
    44  
    45  func (c *CandidateBuffer) emit() {
    46  	prevCandidates := c.clear()
    47  	if len(prevCandidates) > 0 {
    48  		c.onCandidates(prevCandidates)
    49  	}
    50  }
    51  
    52  type OnNextCandidate func(types.RetrievalCandidate)
    53  
    54  // BufferStream consumes a stream of individual candidate results. When a new result comes in, a collection is started, and further results
    55  // are added to the collection until the specified bufferingTime has passed, at which point the collection is passed to the callback setup
    56  // when the Buffer was setup. The timer is reset and the collection emptied until another result comes in. This has the effect of grouping
    57  // results that occur in the same general time frame.
    58  func (c *CandidateBuffer) BufferStream(ctx context.Context, queueCandidates func(context.Context, OnNextCandidate) error, bufferingTime time.Duration) error {
    59  	var wg sync.WaitGroup
    60  	defer func() {
    61  		wg.Wait()
    62  	}()
    63  	err := queueCandidates(ctx, func(candidate types.RetrievalCandidate) {
    64  		var timerCtx context.Context
    65  		c.lk.Lock()
    66  		c.currentCandidates = append(c.currentCandidates, candidate)
    67  		if c.timerCancel != nil {
    68  			c.lk.Unlock()
    69  			return
    70  		}
    71  		timerCtx, c.timerCancel = context.WithCancel(ctx)
    72  		c.lk.Unlock()
    73  		timer := c.clock.Timer(bufferingTime)
    74  		if c.afterEach != nil {
    75  			c.afterEach <- struct{}{}
    76  		}
    77  		wg.Add(1)
    78  		go func() {
    79  			defer wg.Done()
    80  			select {
    81  			case <-timerCtx.Done():
    82  				if !timer.Stop() {
    83  					<-timer.C
    84  				}
    85  			case <-timer.C:
    86  				c.emit()
    87  			}
    88  		}()
    89  	})
    90  	select {
    91  	case <-ctx.Done():
    92  		c.clear()
    93  		return ctx.Err()
    94  	default:
    95  		c.emit()
    96  		return err
    97  	}
    98  }