casket.go

package casket

import (
	"errors"
	"sync"
	"time"
)

var (
	ErrDying      = errors.New("casket: dying")
	ErrStillAlive = errors.New("casket: still alive")
)

type Casket interface {
	Dying() <-chan struct{}
	Dead() <-chan struct{}
	Kill(reason error)
	Err() error
}

type single struct {
	errc   chan error
	dyingc chan struct{}
	deadc  chan struct{}
	err    error
	wg     sync.WaitGroup
}

func (c *single) errorHandler() {
	c.err = <-c.errc
	close(c.dyingc)
	c.wg.Wait()
	close(c.deadc)
}

func (c *single) Dying() <-chan struct{} {
	return c.dyingc
}

func (c *single) Dead() <-chan struct{} {
	return c.deadc
}

func (c *single) Kill(reason error) {
	select {
	case c.errc <- reason:
	case <-c.dyingc:
	}
}

func (c *single) Err() error {
	<-c.dyingc
	return c.err
}

func New() Casket {
	c := &single{
		dyingc: make(chan struct{}),
		deadc:  make(chan struct{}),
		errc:   make(chan error),
	}
	go c.errorHandler()
	return c
}

func Sleep(c Casket, d time.Duration) error {
	select {
	case <-time.After(d):
		return nil
	case <-c.Dying():
		return ErrDying
	}
}

func Do(c Casket, fn func() error) error {
	c.wg.Add(1)
	ch := make(chan error)
	go func() {
		defer c.wg.Done()
		select {
		case ch <- fn():
		case <-c.Dying():
		}
	}()
	select {
	case err := <-ch:
		return err
	case <-c.Dying():
		return ErrDying
	}
}

func Go(c Casket, fn func() error) {
	c.wg.Add(1)
	go func() {
		defer c.wg.Done()
		if err := fn(c); err != ErrDying {
			c.Kill(err)
		}
	}()
}

func Defer(c Casket, fn func()) {
	go func() {
		<-c.Dead()
		fn(err)
	}()
}

// Don't use these
func Dying(c Casket) bool {
	select {
	case <-c.Dying():
		return true
	default:
		return false
	}
}

func Dead(c Casket) bool {
	select {
	case <-c.Dead():
		return true
	default:
		return false
	}
}

func Alive(c Casket) bool {
	return !Dying(c)
}

func Err(c Casket) error {
	if Alive(c) {
		return ErrStillAlive
	} else {
		return c.errc
	}
}