18870: Need to declare NODES as array
[arvados.git] / services / ws / handler.go
index 3d42b9a20892fd929f8d05c1f186ffaa9f24b072..912643ad97c6374006b3fd4b00f90d340157d687 100644 (file)
@@ -1,19 +1,28 @@
-package main
+// Copyright (C) The Arvados Authors. All rights reserved.
+//
+// SPDX-License-Identifier: AGPL-3.0
+
+package ws
 
 import (
-       "encoding/json"
+       "context"
        "io"
+       "sync"
        "time"
 
-       "git.curoverse.com/arvados.git/sdk/go/arvados"
-       log "github.com/Sirupsen/logrus"
+       "git.arvados.org/arvados.git/sdk/go/arvados"
+       "git.arvados.org/arvados.git/sdk/go/stats"
+       "github.com/sirupsen/logrus"
 )
 
 type handler struct {
        Client      arvados.Client
        PingTimeout time.Duration
        QueueSize   int
-       NewSession  func(wsConn, chan<- interface{}) (session, error)
+
+       mtx       sync.Mutex
+       lastDelay map[chan interface{}]stats.Duration
+       setupOnce sync.Once
 }
 
 type handlerStats struct {
@@ -23,119 +32,140 @@ type handlerStats struct {
        EventCount   uint64
 }
 
-func (h *handler) Handle(ws wsConn, incoming <-chan *event) (stats handlerStats) {
-       ctx := contextWithLogger(ws.Request().Context(), log.WithFields(log.Fields{
-               "RemoteAddr": ws.Request().RemoteAddr,
-       }))
+func (h *handler) Handle(ws wsConn, logger logrus.FieldLogger, eventSource eventSource, newSession func(wsConn, chan<- interface{}) (session, error)) (hStats handlerStats) {
+       h.setupOnce.Do(h.setup)
+
+       ctx, cancel := context.WithCancel(ws.Request().Context())
+       defer cancel()
+
+       incoming := eventSource.NewSink()
+       defer incoming.Stop()
 
        queue := make(chan interface{}, h.QueueSize)
-       sess, err := h.NewSession(ws, queue)
-       log := logger(ctx)
+       h.mtx.Lock()
+       h.lastDelay[queue] = 0
+       h.mtx.Unlock()
+       defer func() {
+               h.mtx.Lock()
+               delete(h.lastDelay, queue)
+               h.mtx.Unlock()
+       }()
+
+       sess, err := newSession(ws, queue)
        if err != nil {
-               log.WithError(err).Error("NewSession failed")
+               logger.WithError(err).Error("newSession failed")
                return
        }
 
-       stopped := make(chan struct{})
-       stop := make(chan error, 5)
-
+       // Receive websocket frames from the client and pass them to
+       // sess.Receive().
        go func() {
+               defer cancel()
                buf := make([]byte, 2<<20)
                for {
                        select {
-                       case <-stopped:
+                       case <-ctx.Done():
                                return
                        default:
                        }
                        ws.SetReadDeadline(time.Now().Add(24 * 365 * time.Hour))
                        n, err := ws.Read(buf)
                        buf := buf[:n]
-                       log.WithField("frame", string(buf[:n])).Debug("received frame")
+                       logger.WithField("frame", string(buf[:n])).Debug("received frame")
                        if err == nil && n == cap(buf) {
                                err = errFrameTooBig
                        }
                        if err != nil {
-                               if err != io.EOF {
-                                       log.WithError(err).Info("read error")
+                               if err != io.EOF && ctx.Err() == nil {
+                                       logger.WithError(err).Info("read error")
                                }
-                               stop <- err
                                return
                        }
-                       msg := make(map[string]interface{})
-                       err = json.Unmarshal(buf, &msg)
+                       err = sess.Receive(buf)
                        if err != nil {
-                               log.WithError(err).Info("invalid json from client")
-                               stop <- err
+                               logger.WithError(err).Error("sess.Receive() failed")
                                return
                        }
-                       sess.Receive(msg, buf)
                }
        }()
 
+       // Take items from the outgoing queue, serialize them using
+       // sess.EventMessage() as needed, and send them to the client
+       // as websocket frames.
        go func() {
-               for data := range queue {
+               defer cancel()
+               for {
+                       var ok bool
+                       var data interface{}
+                       select {
+                       case <-ctx.Done():
+                               return
+                       case data, ok = <-queue:
+                               if !ok {
+                                       return
+                               }
+                       }
                        var e *event
                        var buf []byte
                        var err error
-                       log := log
+                       logger := logger
 
                        switch data := data.(type) {
                        case []byte:
                                buf = data
                        case *event:
                                e = data
-                               log = log.WithField("serial", e.Serial)
+                               logger = logger.WithField("serial", e.Serial)
                                buf, err = sess.EventMessage(e)
                                if err != nil {
-                                       log.WithError(err).Error("EventMessage failed")
-                                       stop <- err
-                                       break
+                                       logger.WithError(err).Error("EventMessage failed")
+                                       return
                                } else if len(buf) == 0 {
-                                       log.Debug("skip")
+                                       logger.Debug("skip")
                                        continue
                                }
                        default:
-                               log.WithField("data", data).Error("bad object in client queue")
+                               logger.WithField("data", data).Error("bad object in client queue")
                                continue
                        }
 
-                       log.WithField("frame", string(buf)).Debug("send event")
+                       logger.WithField("frame", string(buf)).Debug("send event")
                        ws.SetWriteDeadline(time.Now().Add(h.PingTimeout))
                        t0 := time.Now()
                        _, err = ws.Write(buf)
                        if err != nil {
-                               log.WithError(err).Error("write failed")
-                               stop <- err
-                               break
+                               if ctx.Err() == nil {
+                                       logger.WithError(err).Error("write failed")
+                               }
+                               return
                        }
-                       log.Debug("sent")
+                       logger.Debug("sent")
 
                        if e != nil {
-                               stats.QueueDelayNs += t0.Sub(e.Received)
+                               hStats.QueueDelayNs += t0.Sub(e.Ready)
+                               h.mtx.Lock()
+                               h.lastDelay[queue] = stats.Duration(time.Since(e.Ready))
+                               h.mtx.Unlock()
                        }
-                       stats.WriteDelayNs += time.Since(t0)
-                       stats.EventBytes += uint64(len(buf))
-                       stats.EventCount++
-               }
-               for _ = range queue {
-                       // Ensure queue can't fill up and block other
-                       // goroutines after we hit a write error.
+                       hStats.WriteDelayNs += time.Since(t0)
+                       hStats.EventBytes += uint64(len(buf))
+                       hStats.EventCount++
                }
        }()
 
        // Filter incoming events against the current subscription
        // list, and forward matching events to the outgoing message
-       // queue. Close the queue and return when the "stopped"
-       // channel closes or the incoming event stream ends. Shut down
-       // the handler if the outgoing queue fills up.
+       // queue. Close the queue and return when the request context
+       // is done/cancelled or the incoming event stream ends. Shut
+       // down the handler if the outgoing queue fills up.
        go func() {
+               defer cancel()
                ticker := time.NewTicker(h.PingTimeout)
                defer ticker.Stop()
 
                for {
                        select {
-                       case <-stopped:
-                               close(queue)
+                       case <-ctx.Done():
                                return
                        case <-ticker.C:
                                // If the outgoing queue is empty,
@@ -144,12 +174,13 @@ func (h *handler) Handle(ws wsConn, incoming <-chan *event) (stats handlerStats)
                                // socket, and prevent an idle socket
                                // from being closed.
                                if len(queue) == 0 {
-                                       queue <- []byte(`{}`)
+                                       select {
+                                       case queue <- []byte(`{}`):
+                                       default:
+                                       }
                                }
-                               continue
-                       case e, ok := <-incoming:
+                       case e, ok := <-incoming.Channel():
                                if !ok {
-                                       close(queue)
                                        return
                                }
                                if !sess.Filter(e) {
@@ -158,14 +189,49 @@ func (h *handler) Handle(ws wsConn, incoming <-chan *event) (stats handlerStats)
                                select {
                                case queue <- e:
                                default:
-                                       stop <- errQueueFull
+                                       logger.WithError(errQueueFull).Error("terminate")
+                                       return
                                }
                        }
                }
        }()
 
-       <-stop
-       close(stopped)
-
+       <-ctx.Done()
        return
 }
+
+func (h *handler) DebugStatus() interface{} {
+       h.mtx.Lock()
+       defer h.mtx.Unlock()
+
+       var s struct {
+               QueueCount    int
+               QueueMin      int
+               QueueMax      int
+               QueueTotal    uint64
+               QueueDelayMin stats.Duration
+               QueueDelayMax stats.Duration
+       }
+       for q, lastDelay := range h.lastDelay {
+               s.QueueCount++
+               n := len(q)
+               s.QueueTotal += uint64(n)
+               if s.QueueMax < n {
+                       s.QueueMax = n
+               }
+               if s.QueueMin > n || s.QueueCount == 1 {
+                       s.QueueMin = n
+               }
+               if (s.QueueDelayMin > lastDelay || s.QueueDelayMin == 0) && lastDelay > 0 {
+                       s.QueueDelayMin = lastDelay
+               }
+               if s.QueueDelayMax < lastDelay {
+                       s.QueueDelayMax = lastDelay
+               }
+       }
+       return &s
+}
+
+func (h *handler) setup() {
+       h.lastDelay = make(map[chan interface{}]stats.Duration)
+}