github.com/ncw/rclone@v1.48.1-0.20190724201158-a35aa1360e3e/fs/sync/pipe.go (about)

     1  package sync
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  
     7  	"github.com/ncw/rclone/fs"
     8  )
     9  
    10  // pipe provides an unbounded channel like experience
    11  //
    12  // Note unlike channels these aren't strictly ordered.
    13  type pipe struct {
    14  	mu        sync.Mutex
    15  	c         chan struct{}
    16  	queue     []fs.ObjectPair
    17  	closed    bool
    18  	totalSize int64
    19  	stats     func(items int, totalSize int64)
    20  }
    21  
    22  func newPipe(stats func(items int, totalSize int64), maxBacklog int) *pipe {
    23  	return &pipe{
    24  		c:     make(chan struct{}, maxBacklog),
    25  		stats: stats,
    26  	}
    27  }
    28  
    29  // Put an pair into the pipe
    30  //
    31  // It returns ok = false if the context was cancelled
    32  //
    33  // It will panic if you call it after Close()
    34  func (p *pipe) Put(ctx context.Context, pair fs.ObjectPair) (ok bool) {
    35  	if ctx.Err() != nil {
    36  		return false
    37  	}
    38  	p.mu.Lock()
    39  	p.queue = append(p.queue, pair)
    40  	size := pair.Src.Size()
    41  	if size > 0 {
    42  		p.totalSize += size
    43  	}
    44  	p.stats(len(p.queue), p.totalSize)
    45  	p.mu.Unlock()
    46  	select {
    47  	case <-ctx.Done():
    48  		return false
    49  	case p.c <- struct{}{}:
    50  	}
    51  	return true
    52  }
    53  
    54  // Get a pair from the pipe
    55  //
    56  // It returns ok = false if the context was cancelled or Close() has
    57  // been called.
    58  func (p *pipe) Get(ctx context.Context) (pair fs.ObjectPair, ok bool) {
    59  	if ctx.Err() != nil {
    60  		return
    61  	}
    62  	select {
    63  	case <-ctx.Done():
    64  		return
    65  	case _, ok = <-p.c:
    66  		if !ok {
    67  			return
    68  		}
    69  	}
    70  	p.mu.Lock()
    71  	pair, p.queue = p.queue[0], p.queue[1:]
    72  	size := pair.Src.Size()
    73  	if size > 0 {
    74  		p.totalSize -= size
    75  	}
    76  	if p.totalSize < 0 {
    77  		p.totalSize = 0
    78  	}
    79  	p.stats(len(p.queue), p.totalSize)
    80  	p.mu.Unlock()
    81  	return pair, true
    82  }
    83  
    84  // Stats reads the number of items in the queue and the totalSize
    85  func (p *pipe) Stats() (items int, totalSize int64) {
    86  	p.mu.Lock()
    87  	items, totalSize = len(p.queue), p.totalSize
    88  	p.mu.Unlock()
    89  	return items, totalSize
    90  }
    91  
    92  // Close the pipe
    93  //
    94  // Writes to a closed pipe will panic as will double closing a pipe
    95  func (p *pipe) Close() {
    96  	p.mu.Lock()
    97  	close(p.c)
    98  	p.closed = true
    99  	p.mu.Unlock()
   100  }