58 lines
877 B
Go
58 lines
877 B
Go
|
package fsm
|
||
|
|
|
||
|
|
import (
|
||
|
|
"sync"
|
||
|
|
"time"
|
||
|
|
)
|
||
|
|
|
||
|
|
type condAndState struct {
|
||
|
|
cond *sync.Cond
|
||
|
|
state *UserState
|
||
|
|
}
|
||
|
|
|
||
|
|
type WaitCallback struct {
|
||
|
|
Waiter sync.Map
|
||
|
|
mu sync.Mutex
|
||
|
|
}
|
||
|
|
|
||
|
|
func NewWaitCallback() *WaitCallback {
|
||
|
|
return &WaitCallback{}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (w *WaitCallback) Done(us *UserState) {
|
||
|
|
if v, ok := w.Waiter.Load(us.UserId); ok {
|
||
|
|
w.mu.Lock()
|
||
|
|
defer w.mu.Unlock()
|
||
|
|
cas := v.(*condAndState)
|
||
|
|
cas.state = us
|
||
|
|
cas.cond.Broadcast()
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (w *WaitCallback) Wait(key int64, timeout time.Duration) *UserState {
|
||
|
|
cas := &condAndState{
|
||
|
|
cond: sync.NewCond(&w.mu),
|
||
|
|
state: nil,
|
||
|
|
}
|
||
|
|
|
||
|
|
v, loaded := w.Waiter.LoadOrStore(key, cas)
|
||
|
|
if loaded {
|
||
|
|
cas = v.(*condAndState)
|
||
|
|
}
|
||
|
|
|
||
|
|
done := make(chan struct{})
|
||
|
|
go func() {
|
||
|
|
w.mu.Lock()
|
||
|
|
defer w.mu.Unlock()
|
||
|
|
cas.cond.Wait()
|
||
|
|
close(done)
|
||
|
|
}()
|
||
|
|
|
||
|
|
select {
|
||
|
|
case <-done:
|
||
|
|
return cas.state
|
||
|
|
case <-time.After(timeout):
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
}
|