[Go] wait until success with timeout

go.mod

module sync/wait

go 1.18

require github.com/stretchr/testify v1.7.1

require (
	github.com/davecgh/go-spew v1.1.0 // indirect
	github.com/pmezard/go-difflib v1.0.0 // indirect
	gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c // indirect
)

go.sum

github.com/davecgh/go-spew v1.1.0 h1:ZDRjVQ15GmhC3fiQ8ni8+OwkZQO4DARzQgrnXU1Liz8=
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
github.com/stretchr/testify v1.7.1 h1:5TQK59W5E3v0r2duFAb7P95B6hEeOyEnHRa8MjYSMTY=
github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c h1:dUUwHk2QECo/6vqA44rthZ8ie2QXMNeKRTHCNY2nXvo=
gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

until.go

package wait

import (
	"context"
	"errors"
	"time"
)

type Exercise = func(ctx context.Context) (done bool, err error)

var ErrTimedOut = errors.New("timed out")

// Until は、処理が成功するまでの一定時間待機する。
//
// * timeout を経過しても exec が成功とならない場合は、ErrTimedOut が返却される
// * exec から エラーが返ったら、その時点で待機が解ける
func Until(ctx context.Context, interval, timeout time.Duration, exec Exercise) error {
	tick := time.NewTicker(interval)
	defer tick.Stop()

	_timeout := make(chan struct{})
	go func() {
		<-time.After(timeout)
		close(_timeout)
	}()

	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-_timeout:
			return ErrTimedOut
		case <-tick.C:
			success, err := exec(ctx)
			if err != nil {
				return err
			}
			if success {
				return nil
			}
		}
	}
}

until_test.go

package wait

import (
	"context"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
)

func TestUntil(t *testing.T) {
	t.Run("block process until exec returns true", func(t *testing.T) {
		t.Parallel()
		// Setup
		var (
			ctx      = context.Background()
			interval = 10 * time.Millisecond
			timeout  = 1 * time.Second

			cnt          = 1
			fn  Exercise = func(_ context.Context) (bool, error) {
				if cnt == 5 {
					return true, nil
				}
				cnt++
				return false, nil
			}
		)
		// Exercise
		err := Until(ctx, interval, timeout, fn)
		assert.NoError(t, err)
	})
	t.Run("timeout duration has passed, then an ErrTimedOut occurs", func(t *testing.T) {
		t.Parallel()
		// Setup
		var (
			ctx               = context.Background()
			interval          = 100 * time.Millisecond
			timeout           = 1 * time.Second
			fn       Exercise = func(_ context.Context) (bool, error) {
				return false, nil // forever unsuccessful.
			}
		)
		// Exercise
		err := Until(ctx, interval, timeout, fn)
		assert.ErrorIs(t, err, ErrTimedOut)
	})
	t.Run("context has canceled, then an context.Deadline occurs", func(t *testing.T) {
		t.Parallel()
		// Setup
		var (
			interval          = 10 * time.Millisecond
			timeout           = 100 * time.Second
			fn       Exercise = func(_ context.Context) (bool, error) {
				return false, nil // forever unsuccessful.
			}
		)
		ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(1*time.Second))
		defer cancel()
		// Exercise
		err := Until(ctx, interval, timeout, fn)
		assert.ErrorIs(t, err, context.DeadlineExceeded)
	})
}