services/ws/handler.go

   1 // Copyright (C) The Arvados Authors. All rights reserved.
   2 //
   3 // SPDX-License-Identifier: AGPL-3.0
   4
   5 package main
   6
   7 import (
   8         "context"
   9         "io"
  10         "sync"
  11         "time"
  12
  13         "git.curoverse.com/arvados.git/sdk/go/arvados"
  14         "git.curoverse.com/arvados.git/sdk/go/stats"
  15 )
  16
  17 type handler struct {
  18         Client      arvados.Client
  19         PingTimeout time.Duration
  20         QueueSize   int
  21
  22         mtx       sync.Mutex
  23         lastDelay map[chan interface{}]stats.Duration
  24         setupOnce sync.Once
  25 }
  26
  27 type handlerStats struct {
  28         QueueDelayNs time.Duration
  29         WriteDelayNs time.Duration
  30         EventBytes   uint64
  31         EventCount   uint64
  32 }
  33
  34 func (h *handler) Handle(ws wsConn, eventSource eventSource, newSession func(wsConn, chan<- interface{}) (session, error)) (hStats handlerStats) {
  35         h.setupOnce.Do(h.setup)
  36
  37         ctx, cancel := context.WithCancel(ws.Request().Context())
  38         defer cancel()
  39         log := logger(ctx)
  40
  41         incoming := eventSource.NewSink()
  42         defer incoming.Stop()
  43
  44         queue := make(chan interface{}, h.QueueSize)
  45         h.mtx.Lock()
  46         h.lastDelay[queue] = 0
  47         h.mtx.Unlock()
  48         defer func() {
  49                 h.mtx.Lock()
  50                 delete(h.lastDelay, queue)
  51                 h.mtx.Unlock()
  52         }()
  53
  54         sess, err := newSession(ws, queue)
  55         if err != nil {
  56                 log.WithError(err).Error("newSession failed")
  57                 return
  58         }
  59
  60         // Receive websocket frames from the client and pass them to
  61         // sess.Receive().
  62         go func() {
  63                 defer cancel()
  64                 buf := make([]byte, 2<<20)
  65                 for {
  66                         select {
  67                         case <-ctx.Done():
  68                                 return
  69                         default:
  70                         }
  71                         ws.SetReadDeadline(time.Now().Add(24 * 365 * time.Hour))
  72                         n, err := ws.Read(buf)
  73                         buf := buf[:n]
  74                         log.WithField("frame", string(buf[:n])).Debug("received frame")
  75                         if err == nil && n == cap(buf) {
  76                                 err = errFrameTooBig
  77                         }
  78                         if err != nil {
  79                                 if err != io.EOF && ctx.Err() == nil {
  80                                         log.WithError(err).Info("read error")
  81                                 }
  82                                 return
  83                         }
  84                         err = sess.Receive(buf)
  85                         if err != nil {
  86                                 log.WithError(err).Error("sess.Receive() failed")
  87                                 return
  88                         }
  89                 }
  90         }()
  91
  92         // Take items from the outgoing queue, serialize them using
  93         // sess.EventMessage() as needed, and send them to the client
  94         // as websocket frames.
  95         go func() {
  96                 defer cancel()
  97                 for {
  98                         var ok bool
  99                         var data interface{}
 100                         select {
 101                         case <-ctx.Done():
 102                                 return
 103                         case data, ok = <-queue:
 104                                 if !ok {
 105                                         return
 106                                 }
 107                         }
 108                         var e *event
 109                         var buf []byte
 110                         var err error
 111                         log := log
 112
 113                         switch data := data.(type) {
 114                         case []byte:
 115                                 buf = data
 116                         case *event:
 117                                 e = data
 118                                 log = log.WithField("serial", e.Serial)
 119                                 buf, err = sess.EventMessage(e)
 120                                 if err != nil {
 121                                         log.WithError(err).Error("EventMessage failed")
 122                                         return
 123                                 } else if len(buf) == 0 {
 124                                         log.Debug("skip")
 125                                         continue
 126                                 }
 127                         default:
 128                                 log.WithField("data", data).Error("bad object in client queue")
 129                                 continue
 130                         }
 131
 132                         log.WithField("frame", string(buf)).Debug("send event")
 133                         ws.SetWriteDeadline(time.Now().Add(h.PingTimeout))
 134                         t0 := time.Now()
 135                         _, err = ws.Write(buf)
 136                         if err != nil {
 137                                 if ctx.Err() == nil {
 138                                         log.WithError(err).Error("write failed")
 139                                 }
 140                                 return
 141                         }
 142                         log.Debug("sent")
 143
 144                         if e != nil {
 145                                 hStats.QueueDelayNs += t0.Sub(e.Ready)
 146                                 h.mtx.Lock()
 147                                 h.lastDelay[queue] = stats.Duration(time.Since(e.Ready))
 148                                 h.mtx.Unlock()
 149                         }
 150                         hStats.WriteDelayNs += time.Since(t0)
 151                         hStats.EventBytes += uint64(len(buf))
 152                         hStats.EventCount++
 153                 }
 154         }()
 155
 156         // Filter incoming events against the current subscription
 157         // list, and forward matching events to the outgoing message
 158         // queue. Close the queue and return when the request context
 159         // is done/cancelled or the incoming event stream ends. Shut
 160         // down the handler if the outgoing queue fills up.
 161         go func() {
 162                 defer cancel()
 163                 ticker := time.NewTicker(h.PingTimeout)
 164                 defer ticker.Stop()
 165
 166                 for {
 167                         select {
 168                         case <-ctx.Done():
 169                                 return
 170                         case <-ticker.C:
 171                                 // If the outgoing queue is empty,
 172                                 // send an empty message. This can
 173                                 // help detect a disconnected network
 174                                 // socket, and prevent an idle socket
 175                                 // from being closed.
 176                                 if len(queue) == 0 {
 177                                         select {
 178                                         case queue <- []byte(`{}`):
 179                                         default:
 180                                         }
 181                                 }
 182                         case e, ok := <-incoming.Channel():
 183                                 if !ok {
 184                                         return
 185                                 }
 186                                 if !sess.Filter(e) {
 187                                         continue
 188                                 }
 189                                 select {
 190                                 case queue <- e:
 191                                 default:
 192                                         log.WithError(errQueueFull).Error("terminate")
 193                                         return
 194                                 }
 195                         }
 196                 }
 197         }()
 198
 199         <-ctx.Done()
 200         return
 201 }
 202
 203 func (h *handler) DebugStatus() interface{} {
 204         h.mtx.Lock()
 205         defer h.mtx.Unlock()
 206
 207         var s struct {
 208                 QueueCount    int
 209                 QueueMin      int
 210                 QueueMax      int
 211                 QueueTotal    uint64
 212                 QueueDelayMin stats.Duration
 213                 QueueDelayMax stats.Duration
 214         }
 215         for q, lastDelay := range h.lastDelay {
 216                 s.QueueCount++
 217                 n := len(q)
 218                 s.QueueTotal += uint64(n)
 219                 if s.QueueMax < n {
 220                         s.QueueMax = n
 221                 }
 222                 if s.QueueMin > n || s.QueueCount == 1 {
 223                         s.QueueMin = n
 224                 }
 225                 if (s.QueueDelayMin > lastDelay || s.QueueDelayMin == 0) && lastDelay > 0 {
 226                         s.QueueDelayMin = lastDelay
 227                 }
 228                 if s.QueueDelayMax < lastDelay {
 229                         s.QueueDelayMax = lastDelay
 230                 }
 231         }
 232         return &s
 233 }
 234
 235 func (h *handler) setup() {
 236         h.lastDelay = make(map[chan interface{}]stats.Duration)
 237 }