counter.go

package todocounter

import (
	"sync"
)

// Counter records things remaining to process. It is needed for complicated
// cases where multiple goroutines are spawned to process items, and they may
// generate more items to process. For example, say a query over a set of nodes
// may yield either a result value, or more nodes to query. Signaling is subtly
// complicated, because the queue may be empty while items are being processed,
// that will end up adding more items to the queue.
//
// Use Counter like this:
//
//	todos := make(chan int, 10)
//	ctr := todoctr.NewCounter()
//
//	process := func(item int) {
//	  fmt.Println("processing %d\n...", item)
//
//	  // this task may randomly generate more tasks
//	  if rand.Intn(5) == 0 {
//	    todos<- item + 1
//	    ctr.Increment(1) // increment counter for new task.
//	  }
//
//	  ctr.Decrement(1) // decrement one to signal the task being done.
//	}
//
//	// add some tasks.
//	todos<- 1
//	todos<- 2
//	todos<- 3
//	todos<- 4
//	ctr.Increment(4)
//
//	for {
//	  select {
//	  case item := <- todos:
//	    go process(item)
//	  case <-ctr.Done():
//	    fmt.Println("done processing everything.")
//	    close(todos)
//	  }
//	}
type Counter interface {
	// Incrememnt adds a number of todos to track.
	// If the counter is **below** zero, it panics.
	Increment(i uint32)

	// Decrement removes a number of todos to track.
	// If the count drops to zero, signals done and destroys the counter.
	// If the count drops **below** zero, panics. It means you have tried to remove
	// more things than you added, i.e. sync issues.
	Decrement(i uint32)

	// Done returns a channel to wait upon. Use it in selects:
	//
	//  select {
	//  case <-ctr.Done():
	//    // done processing all items
	//  }
	//
	Done() <-chan struct{}

	Remaining() int32
}

type todoCounter struct {
	count int32
	done  chan struct{}
	sync.RWMutex
}

// NewSyncCounter constructs a new counter
func NewSyncCounter() Counter {
	return &todoCounter{
		done: make(chan struct{}),
	}
}

func (c *todoCounter) Increment(i uint32) {
	c.Lock()
	defer c.Unlock()

	if c.count < 0 {
		panic("counter already signaled done. use a new counter.")
	}

	// increment count
	c.count += int32(i)
}

// Decrement removes a number of todos to track.
// If the count drops to zero, signals done and destroys the counter.
// If the count drops **below** zero, panics. It means you have tried to remove
// more things than you added, i.e. sync issues.
func (c *todoCounter) Decrement(i uint32) {
	c.Lock()
	defer c.Unlock()

	if c.count < 0 {
		panic("counter already signaled done. probably have sync issues.")
	}

	if int32(i) > c.count {
		panic("decrement amount creater than counter. sync issues.")
	}

	c.count -= int32(i)
	if c.count == 0 { // done! signal it.
		c.count--     // set it to -1 to prevent reuse
		close(c.done) // a closed channel will always return nil
	}
}

func (c *todoCounter) Done() <-chan struct{} {
	return c.done
}

func (c *todoCounter) Remaining() int32 {
	c.RLock()
	defer c.RUnlock()

	return c.count
}