JCC-CSScheduler/manager/internal/jobmgr/event_set.go

73 lines
1.3 KiB
Go

package jobmgr
import (
"context"
"sync"
"gitlink.org.cn/cloudream/common/pkgs/future"
"gitlink.org.cn/cloudream/common/utils/lo2"
)
type EventWaitCondition func(evt Event) bool
type EventWaiter struct {
condition EventWaitCondition
future *future.SetValueFuture[Event]
}
type EventSet struct {
events []Event
waiters []EventWaiter
lock sync.Mutex
}
func NewEventSet() EventSet {
return EventSet{}
}
func (s *EventSet) Post(evt Event) {
s.lock.Lock()
defer s.lock.Unlock()
// 一个事件能唤醒多个等待者
used := false
for i, waiter := range s.waiters {
if waiter.condition(evt) {
s.waiters = lo2.RemoveAt(s.waiters, i)
waiter.future.SetValue(evt)
used = true
}
}
if !used {
s.events = append(s.events, evt)
}
}
func (s *EventSet) Wait(ctx context.Context, cond EventWaitCondition) (Event, bool) {
s.lock.Lock()
defer s.lock.Unlock()
// 一个等待者只能等待一个事件
for i, evt := range s.events {
if cond(evt) {
s.events = lo2.RemoveAt(s.events, i)
return evt, true
}
}
fut := future.NewSetValue[Event]()
waiter := EventWaiter{
condition: cond,
future: fut,
}
s.events = append(s.events, waiter)
val, err := fut.WaitValue(ctx)
if err != nil {
return nil, false
}
return val, true
}