github.com/coyove/common@v0.0.0-20240403014525-f70e643f9de8/waitobject/wait.go (about)

     1  package waitobject
     2  
     3  import (
     4  	"fmt"
     5  	"sync"
     6  	"time"
     7  )
     8  
     9  func debugprint(v ...interface{}) {
    10  	if debug {
    11  		fmt.Println(time.Now().Format(time.RFC3339), fmt.Sprint(v...))
    12  	}
    13  }
    14  
    15  type Object struct {
    16  	mu      sync.Mutex
    17  	v       interface{}
    18  	sig     *sync.Cond
    19  	rev     *notifier
    20  	touched bool
    21  }
    22  
    23  func New() *Object {
    24  	o := &Object{}
    25  	o.sig = sync.NewCond(&o.mu)
    26  	return o
    27  }
    28  
    29  func (o *Object) Touch(f func(old interface{}) interface{}) {
    30  	o.mu.Lock()
    31  	debugprint("broadcast by touching")
    32  	o.v = f(o.v)
    33  	o.touched = true
    34  	o.sig.Signal()
    35  	o.mu.Unlock()
    36  }
    37  
    38  func (o *Object) SetValue(f func(v interface{}) interface{}) interface{} {
    39  	o.mu.Lock()
    40  	defer o.mu.Unlock()
    41  	old := o.v
    42  	if f != nil {
    43  		o.touched = false
    44  		o.v = f(o.v)
    45  	}
    46  	return old
    47  }
    48  
    49  // SetWaitDeadline sets the deadline of Wait(), note that its precision is 1s
    50  func (o *Object) SetWaitDeadline(t time.Time) {
    51  	o.mu.Lock()
    52  	defer o.mu.Unlock()
    53  
    54  	if o.rev != nil {
    55  		// Current object has a notifier in the timeoutwheel
    56  		// invalidate to prevent it from firing any old timeout events in the future
    57  		o.rev.invalidate()
    58  		o.rev = nil
    59  	}
    60  
    61  	if t.IsZero() || t == Eternal {
    62  		// Caller wants to cancel the deadline
    63  		return
    64  	}
    65  
    66  	o.rev = &notifier{deadline: t.Unix(), obj: o}
    67  	if o.isTimedout() {
    68  		debugprint("direct (already) timeout")
    69  		o.sig.Broadcast()
    70  		return
    71  	}
    72  
    73  	ts := &timeoutWheel.secmin[t.Second()][t.Minute()]
    74  	ts.Lock()
    75  	ts.list = append(ts.list, o.rev)
    76  	ts.Unlock()
    77  }
    78  
    79  func (o *Object) IsTimedout() bool {
    80  	o.mu.Lock()
    81  	defer o.mu.Unlock()
    82  	return o.isTimedout()
    83  }
    84  
    85  func (o *Object) isTimedout() bool {
    86  	if o.rev == nil {
    87  		return false
    88  	}
    89  
    90  	now := time.Now().Unix()
    91  	out := now >= o.rev.deadline
    92  
    93  	debugprint("isTimedout: ", out, ", now: ", now, ", deadline: ", o.rev.deadline)
    94  	return out
    95  }
    96  
    97  func (o *Object) Wait() (interface{}, bool) {
    98  	o.mu.Lock()
    99  	defer o.mu.Unlock()
   100  
   101  	// Before waiting for any data, return early if it is already timed out
   102  	if o.isTimedout() {
   103  		return nil, false
   104  	}
   105  
   106  	if o.touched {
   107  		o.touched = false
   108  		return o.v, true
   109  	}
   110  
   111  	debugprint("wait start, v: ", o.v)
   112  	o.sig.Wait()
   113  	o.touched = false
   114  	debugprint("wait end, v: ", o.v)
   115  
   116  	// After receiving any data, return early if it is already timed out
   117  	if o.isTimedout() {
   118  		return nil, false
   119  	}
   120  
   121  	return o.v, true
   122  }