Move keepalive out of mgm pkg

keepalive/keep_alive.go

@@ -0,0 +1,152 @@
+package keepalive
+
+import (
+	"context"
+	"sync"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/metadata"
+)
+
+const (
+	reversProxyHeaderKey = "x-netbird-peer"
+	grpcVersionHeaderKey = "version"
+	keepAliveInterval    = 30 * time.Second
+)
+
+type KeepAlive struct {
+	sync.RWMutex
+	ticker       *time.Ticker
+	done         chan struct{}
+	streams      map[string]*ioMonitor
+	keepAliveMsg interface{}
+}
+
+func NewKeepAlive(keepAliveMsg interface{}) *KeepAlive {
+	ka := &KeepAlive{
+		ticker:       time.NewTicker(1 * time.Second),
+		done:         make(chan struct{}),
+		streams:      make(map[string]*ioMonitor),
+		keepAliveMsg: keepAliveMsg,
+	}
+	go ka.start()
+	return ka
+}
+
+func (k *KeepAlive) StreamInterceptor() grpc.StreamServerInterceptor {
+	return func(srv interface{}, stream grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
+		address, supported := k.keepAliveIsSupported(stream.Context())
+		if !supported {
+			return handler(srv, stream)
+		}
+
+		m := &ioMonitor{
+			sync.Mutex{},
+			stream,
+			time.Now(),
+		}
+
+		k.addIoMonitor(address, m)
+
+		return handler(srv, m)
+	}
+}
+
+func (k *KeepAlive) UnaryInterceptor() grpc.UnaryServerInterceptor {
+	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp interface{}, err error) {
+		address, supported := k.keepAliveIsSupported(ctx)
+		if supported {
+			k.updateLastSeen(address)
+		}
+		return handler(ctx, req)
+	}
+}
+
+func (k *KeepAlive) Stop() {
+	select {
+	case k.done <- struct{}{}:
+		k.ticker.Stop()
+		return
+	default:
+	}
+}
+
+func (k *KeepAlive) start() {
+	for {
+		select {
+		case <-k.done:
+			return
+		case t := <-k.ticker.C:
+			k.checkKeepAlive(t)
+		}
+	}
+}
+
+func (k *KeepAlive) checkKeepAlive(now time.Time) {
+	k.Lock()
+	defer k.Unlock()
+	for addr, m := range k.streams {
+		if k.isKeepAliveOutDated(now, m) {
+			continue
+		}
+		err := k.sendKeepAlive(m)
+		if err != nil {
+			log.Debugf("stop keepalive for: %s", addr)
+			delete(k.streams, addr)
+		}
+	}
+}
+
+func (k *KeepAlive) keepAliveIsSupported(ctx context.Context) (string, bool) {
+	md, ok := metadata.FromIncomingContext(ctx)
+	if !ok {
+		log.Warnf("metadata not found")
+		return "", false
+	}
+
+	peerAddress := k.addressFromHeader(md)
+	if peerAddress == "" {
+		log.Debugf("peer is not using reverse proxy")
+		return "", false
+	}
+
+	if len(md.Get(grpcVersionHeaderKey)) == 0 {
+		log.Debugf("version info not found")
+		return "", false
+	}
+	return peerAddress, true
+}
+
+func (k *KeepAlive) addIoMonitor(address string, m *ioMonitor) {
+	k.Lock()
+	defer k.Unlock()
+	k.streams[address] = m
+}
+
+func (k *KeepAlive) sendKeepAlive(m *ioMonitor) error {
+	return m.sendMsg(k.keepAliveMsg)
+}
+
+func (k *KeepAlive) updateLastSeen(address string) {
+	k.RLock()
+	m, ok := k.streams[address]
+	k.RUnlock()
+	if !ok {
+		return
+	}
+	m.updateLastSeen()
+}
+
+func (k *KeepAlive) addressFromHeader(md metadata.MD) string {
+	peer := md.Get(reversProxyHeaderKey)
+	if len(peer) == 0 {
+		return ""
+	}
+	return peer[0]
+}
+
+func (k *KeepAlive) isKeepAliveOutDated(now time.Time, m *ioMonitor) bool {
+	return now.Sub(m.getLastSeen()) < keepAliveInterval
+}

keepalive/monitor.go

@@ -0,0 +1,32 @@
+package keepalive
+
+import (
+	"sync"
+	"time"
+
+	"google.golang.org/grpc"
+)
+
+type ioMonitor struct {
+	mu sync.Mutex
+	grpc.ServerStream
+	lastSeen time.Time
+}
+
+func (l *ioMonitor) sendMsg(m interface{}) error {
+	l.updateLastSeen()
+	return l.ServerStream.SendMsg(m)
+}
+
+func (l *ioMonitor) updateLastSeen() {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+	l.lastSeen = time.Now()
+}
+
+func (l *ioMonitor) getLastSeen() time.Time {
+	l.mu.Lock()
+	t := l.lastSeen
+	l.mu.Unlock()
+	return t
+}