sema.go

sema.go

package jdksync

import (
	"sync"
)

type Sema struct {
	s chan chan struct{}

	mu      sync.Mutex
	done    chan struct{}
	waiting []chan struct{}
}

func NewSema() *Sema {
	return &Sema{
		// Needs buffer=1 so that the Block can put something on and then return,
		// guaranteeing that Block, Unblock don't race. (when Block returns, you
		// are guaranteed to actually be blocking)
		s: make(chan chan struct{}, 1),

		mu:   sync.Mutex{},
		done: make(chan struct{}),
	}
}

func (s *Sema) Block() {
	once := make(chan struct{})
	onceOnce := sync.Once{}
	done := make(chan struct{})

	s.mu.Lock()
	s.done = done
	s.mu.Unlock()

	go func() {
		for {
			next := make(chan struct{})

			s.mu.Lock()
			s.waiting = append(s.waiting, next)
			s.mu.Unlock()

			select {
			case s.s <- next:
				onceOnce.Do(func() {
					close(once)
				})
			case <-done:
				return
			}
		}
	}()

	select {
	case <-once:
	}
}

func (s *Sema) Unblock() {
	s.mu.Lock()
	close(s.done)
	for _, w := range s.waiting {
		close(w)
	}
	s.waiting = []chan struct{}{}
	s.mu.Unlock()
}

func (s *Sema) Wait() <-chan struct{} {
	select {
	case w := <-s.s:
		return w
	default:
	}

	w := make(chan struct{})
	close(w)
	return w
}