github.com/gdamore/mangos@v1.4.0/waiter.go (about)

     1  // Copyright 2015 The Mangos Authors
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use file except in compliance with the License.
     5  // You may obtain a copy of the license at
     6  //
     7  //    http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package mangos
    16  
    17  import (
    18  	"sync"
    19  	"time"
    20  )
    21  
    22  // CondTimed is a condition variable (ala sync.Cond) but inclues a timeout.
    23  type CondTimed struct {
    24  	sync.Cond
    25  }
    26  
    27  // WaitRelTimeout is like Wait, but it times out.  The fact that
    28  // it timed out can be determined by checking the return value.  True
    29  // indicates that it woke up without a timeout (signaled another way),
    30  // whereas false indicates a timeout occurred.
    31  func (cv *CondTimed) WaitRelTimeout(when time.Duration) bool {
    32  	timer := time.AfterFunc(when, func() {
    33  		cv.L.Lock()
    34  		cv.Broadcast()
    35  		cv.L.Unlock()
    36  	})
    37  	cv.Wait()
    38  	return timer.Stop()
    39  }
    40  
    41  // WaitAbsTimeout is like WaitRelTimeout, but expires on an absolute time
    42  // instead of a relative one.
    43  func (cv *CondTimed) WaitAbsTimeout(when time.Time) bool {
    44  	now := time.Now()
    45  	if when.After(now) {
    46  		return cv.WaitRelTimeout(when.Sub(now))
    47  	}
    48  	return cv.WaitRelTimeout(time.Duration(0))
    49  }
    50  
    51  // Waiter is a way to wait for completion, but it includes a timeout.  It
    52  // is similar in some respects to sync.WaitGroup.
    53  type Waiter struct {
    54  	cv  CondTimed
    55  	cnt int
    56  	sync.Mutex
    57  }
    58  
    59  // Init must be called to initialize the Waiter.
    60  func (w *Waiter) Init() {
    61  	w.cv.L = w
    62  	w.cnt = 0
    63  }
    64  
    65  // Add adds a new go routine/item to wait for. This should be called before
    66  // starting go routines you want to wait for, for example.
    67  func (w *Waiter) Add() {
    68  	w.Lock()
    69  	w.cnt++
    70  	w.Unlock()
    71  }
    72  
    73  // Done is called when the item to wait for is done. There should be a one to
    74  // one correspondance between Add and Done.  When the count drops to zero,
    75  // any callers blocked in Wait() are woken.  If the count drops below zero,
    76  // it panics.
    77  func (w *Waiter) Done() {
    78  	w.Lock()
    79  	w.cnt--
    80  	if w.cnt < 0 {
    81  		panic("wait count dropped < 0")
    82  	}
    83  	if w.cnt == 0 {
    84  		w.cv.Broadcast()
    85  	}
    86  	w.Unlock()
    87  }
    88  
    89  // Wait waits without a timeout.  It only completes when the count drops
    90  // to zero.
    91  func (w *Waiter) Wait() {
    92  	w.Lock()
    93  	for w.cnt != 0 {
    94  		w.cv.Wait()
    95  	}
    96  	w.Unlock()
    97  }
    98  
    99  // WaitRelTimeout waits until either the count drops to zero, or the timeout
   100  // expires.  It returns true if the count is zero, false otherwise.
   101  func (w *Waiter) WaitRelTimeout(d time.Duration) bool {
   102  	w.Lock()
   103  	for w.cnt != 0 {
   104  		if !w.cv.WaitRelTimeout(d) {
   105  			break
   106  		}
   107  	}
   108  	done := w.cnt == 0
   109  	w.Unlock()
   110  	return done
   111  }
   112  
   113  // WaitAbsTimeout is like WaitRelTimeout, but waits until an absolute time.
   114  func (w *Waiter) WaitAbsTimeout(t time.Time) bool {
   115  	w.Lock()
   116  	for w.cnt != 0 {
   117  		if !w.cv.WaitAbsTimeout(t) {
   118  			break
   119  		}
   120  	}
   121  	done := w.cnt == 0
   122  	w.Unlock()
   123  	return done
   124  }