Implement WaitTimeout func

pkg/token/token.go

@@ -2,6 +2,8 @@ package token
 
 import (
 	"fmt"
+	"sync"
+	"sync/atomic"
 	"time"
 
 	mqtt "github.com/eclipse/paho.mqtt.golang"
@@ -37,9 +39,34 @@ func (ts *Tokens) Wait() bool {
 // Token to complete, returns true if it returned before the timeout or
 // returns false if the timeout occurred. In the case of a timeout the Token
 // does not have an error set in case the caller wishes to wait again.
-func (ts *Tokens) WaitTimeout(time.Duration) bool {
+func (ts *Tokens) WaitTimeout(d time.Duration) bool {
+	result := atomic.Value{}
+	result.Store(true)
+
+	wg := &sync.WaitGroup{}
+	wg.Add(len(ts.tokens))
+
+	for _, t := range ts.tokens {
+		t := t
+		go func() {
+			defer wg.Done()
+			result.Store(result.Load().(bool) && t.WaitTimeout(d))
+		}()
+	}
+
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		wg.Wait()
+	}()
+
+	select {
+	case <-done:
+		return result.Load().(bool)
+	case <-time.After(d):
 		return false
 	}
+}
 
 // Done returns a channel that is closed when the flow associated
 // with the Token completes. Clients should call Error after the