github.com/zignig/go-ipfs@v0.0.0-20141111235910-c9e5fdf55a52/util/todocounter/counter.go (about)

     1  package todocounter
     2  
     3  import (
     4  	"sync"
     5  )
     6  
     7  // Counter records things remaining to process. It is needed for complicated
     8  // cases where multiple goroutines are spawned to process items, and they may
     9  // generate more items to process. For example, say a query over a set of nodes
    10  // may yield either a result value, or more nodes to query. Signaling is subtly
    11  // complicated, because the queue may be empty while items are being processed,
    12  // that will end up adding more items to the queue.
    13  //
    14  // Use Counter like this:
    15  //
    16  //    todos := make(chan int, 10)
    17  //    ctr := todoctr.NewCounter()
    18  //
    19  //    process := func(item int) {
    20  //      fmt.Println("processing %d\n...", item)
    21  //
    22  //      // this task may randomly generate more tasks
    23  //      if rand.Intn(5) == 0 {
    24  //        todos<- item + 1
    25  //        ctr.Increment(1) // increment counter for new task.
    26  //      }
    27  //
    28  //      ctr.Decrement(1) // decrement one to signal the task being done.
    29  //    }
    30  //
    31  //    // add some tasks.
    32  //    todos<- 1
    33  //    todos<- 2
    34  //    todos<- 3
    35  //    todos<- 4
    36  //    ctr.Increment(4)
    37  //
    38  //    for {
    39  //      select {
    40  //      case item := <- todos:
    41  //        go process(item)
    42  //      case <-ctr.Done():
    43  //        fmt.Println("done processing everything.")
    44  //        close(todos)
    45  //      }
    46  //    }
    47  type Counter interface {
    48  	// Incrememnt adds a number of todos to track.
    49  	// If the counter is **below** zero, it panics.
    50  	Increment(i uint32)
    51  
    52  	// Decrement removes a number of todos to track.
    53  	// If the count drops to zero, signals done and destroys the counter.
    54  	// If the count drops **below** zero, panics. It means you have tried to remove
    55  	// more things than you added, i.e. sync issues.
    56  	Decrement(i uint32)
    57  
    58  	// Done returns a channel to wait upon. Use it in selects:
    59  	//
    60  	//  select {
    61  	//  case <-ctr.Done():
    62  	//    // done processing all items
    63  	//  }
    64  	//
    65  	Done() <-chan struct{}
    66  }
    67  
    68  type todoCounter struct {
    69  	count int32
    70  	done  chan struct{}
    71  	sync.RWMutex
    72  }
    73  
    74  // NewSyncCounter constructs a new counter
    75  func NewSyncCounter() Counter {
    76  	return &todoCounter{
    77  		done: make(chan struct{}),
    78  	}
    79  }
    80  
    81  func (c *todoCounter) Increment(i uint32) {
    82  	c.Lock()
    83  	defer c.Unlock()
    84  
    85  	if c.count < 0 {
    86  		panic("counter already signaled done. use a new counter.")
    87  	}
    88  
    89  	// increment count
    90  	c.count += int32(i)
    91  }
    92  
    93  // Decrement removes a number of todos to track.
    94  // If the count drops to zero, signals done and destroys the counter.
    95  // If the count drops **below** zero, panics. It means you have tried to remove
    96  // more things than you added, i.e. sync issues.
    97  func (c *todoCounter) Decrement(i uint32) {
    98  	c.Lock()
    99  	defer c.Unlock()
   100  
   101  	if c.count < 0 {
   102  		panic("counter already signaled done. probably have sync issues.")
   103  	}
   104  
   105  	if int32(i) > c.count {
   106  		panic("decrement amount creater than counter. sync issues.")
   107  	}
   108  
   109  	c.count -= int32(i)
   110  	if c.count == 0 { // done! signal it.
   111  		c.count--     // set it to -1 to prevent reuse
   112  		close(c.done) // a closed channel will always return nil
   113  	}
   114  }
   115  
   116  func (c *todoCounter) Done() <-chan struct{} {
   117  	return c.done
   118  }