k3s/vendor/github.com/rancher/kine/pkg/server/watch.go
Darren Shepherd a005219c5b Update vendor
2019-08-22 09:14:47 -07:00

144 lines
2.8 KiB
Go

package server
import (
"context"
"sync"
"sync/atomic"
"github.com/coreos/etcd/etcdserver/etcdserverpb"
"github.com/coreos/etcd/mvcc/mvccpb"
"github.com/sirupsen/logrus"
)
var (
watchID int64
)
func (s *KVServerBridge) Watch(ws etcdserverpb.Watch_WatchServer) error {
w := watcher{
server: ws,
backend: s.limited.backend,
watches: map[int64]func(){},
}
defer w.Close()
for {
msg, err := ws.Recv()
if err != nil {
return err
}
if msg.GetCreateRequest() != nil {
w.Start(msg.GetCreateRequest())
} else if msg.GetCancelRequest() != nil {
w.Cancel(msg.GetCancelRequest().WatchId)
}
}
}
type watcher struct {
sync.Mutex
wg sync.WaitGroup
backend Backend
server etcdserverpb.Watch_WatchServer
watches map[int64]func()
}
func (w *watcher) Start(r *etcdserverpb.WatchCreateRequest) {
w.Lock()
defer w.Unlock()
ctx, cancel := context.WithCancel(context.Background())
id := atomic.AddInt64(&watchID, 1)
w.watches[id] = cancel
w.wg.Add(1)
key := string(r.Key)
logrus.Debugf("WATCH START id=%d, key=%s, revision=%d", id, key, r.StartRevision)
go func() {
defer w.wg.Done()
if err := w.server.Send(&etcdserverpb.WatchResponse{
Header: &etcdserverpb.ResponseHeader{},
Created: true,
WatchId: id,
}); err != nil {
w.Cancel(id)
return
}
for events := range w.backend.Watch(ctx, key, r.StartRevision) {
if len(events) == 0 {
continue
}
if logrus.IsLevelEnabled(logrus.DebugLevel) {
for _, event := range events {
logrus.Debugf("WATCH READ id=%d, key=%s, revision=%d", id, event.KV.Key, event.KV.ModRevision)
}
}
if err := w.server.Send(&etcdserverpb.WatchResponse{
Header: txnHeader(events[len(events)-1].KV.ModRevision),
WatchId: id,
Events: toEvents(events...),
}); err != nil {
w.Cancel(id)
return
}
}
logrus.Debugf("WATCH CLOSE id=%d, key=%s", id, key)
}()
}
func toEvents(events ...*Event) []*mvccpb.Event {
ret := make([]*mvccpb.Event, 0, len(events))
for _, e := range events {
ret = append(ret, toEvent(e))
}
return ret
}
func toEvent(event *Event) *mvccpb.Event {
e := &mvccpb.Event{
Kv: toKV(event.KV),
PrevKv: toKV(event.PrevKV),
}
if event.Delete {
e.Type = mvccpb.DELETE
} else {
e.Type = mvccpb.PUT
}
return e
}
func (w *watcher) Cancel(watchID int64) {
w.Lock()
defer w.Unlock()
if cancel, ok := w.watches[watchID]; ok {
cancel()
delete(w.watches, watchID)
}
err := w.server.Send(&etcdserverpb.WatchResponse{
Header: &etcdserverpb.ResponseHeader{},
Canceled: true,
WatchId: watchID,
})
if err != nil {
logrus.Errorf("Failed to send cancel response for watchID %d: %v", watchID, err)
}
}
func (w *watcher) Close() {
w.Lock()
defer w.Unlock()
for _, v := range w.watches {
v()
}
w.wg.Wait()
}