Added monitor rate limiting

Author: cmaglie

commands/daemon/monitor.go

@@ -19,6 +19,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"sync/atomic"
 
 	"github.com/arduino/arduino-cli/arduino/monitors"
 	rpc "github.com/arduino/arduino-cli/rpc/monitor"
@@ -79,6 +80,17 @@ func (s *MonitorService) StreamingOpen(stream rpc.Monitor_StreamingOpenServer) e
 	streamClosed := make(chan error)
 	targetClosed := make(chan error)
 
+	// set rate limiting window
+	bufferSize := int(config.GetRecvRateLimitBuffer())
+	rateLimitEnabled := (bufferSize > 0)
+	if !rateLimitEnabled {
+		bufferSize = 1024
+	}
+	buffer := make([]byte, bufferSize)
+	bufferUsed := 0
+
+	var writeSlots int32
+
 	// now we can read the other messages and re-route to the monitor...
 	go func() {
 		for {
@@ -95,6 +107,11 @@ func (s *MonitorService) StreamingOpen(stream rpc.Monitor_StreamingOpenServer) e
 				break
 			}
 
+			if rateLimitEnabled {
+				// Increase rate limiter write slots
+				atomic.AddInt32(&writeSlots, msg.GetRecvAcknowledge())
+			}
+
 			if _, err := mon.Write(msg.GetData()); err != nil {
 				// error writing to target
 				targetClosed <- err
@@ -105,27 +122,59 @@ func (s *MonitorService) StreamingOpen(stream rpc.Monitor_StreamingOpenServer) e
 
 	// ...and read from the monitor and forward to the output stream
 	go func() {
-		buf := make([]byte, 8)
+		dropBuffer := make([]byte, 10240)
+		dropped := 0
 		for {
-			n, err := mon.Read(buf)
-			if err != nil {
-				// error reading from target
-				targetClosed <- err
-				break
-			}
-
-			if n == 0 {
-				// target was closed
-				targetClosed <- nil
-				break
+			if bufferUsed < bufferSize {
+				if n, err := mon.Read(buffer[bufferUsed:]); err != nil {
+					// error reading from target
+					targetClosed <- err
+					break
+				} else if n == 0 {
+					// target was closed
+					targetClosed <- nil
+					break
+				} else {
+					bufferUsed += n
+				}
+			} else {
+				// FIXME: a very rare condition but still...
+				// we may be waiting here while, in the meantime, a transmit slot is
+				// freed: in this case the (filled) buffer will stay in the server
+				// until the following Read exits (-> the next char arrives from the
+				// monitor).
+
+				if n, err := mon.Read(dropBuffer); err != nil {
+					// error reading from target
+					targetClosed <- err
+					break
+				} else if n == 0 {
+					// target was closed
+					targetClosed <- nil
+					break
+				} else {
+					dropped += n
+				}
 			}
 
-			if err = stream.Send(&rpc.StreamingOpenResp{
-				Data: buf[:n],
-			}); err != nil {
-				// error sending to stream
-				streamClosed <- err
-				break
+			slots := atomic.LoadInt32(&writeSlots)
+			if !rateLimitEnabled || slots > 0 {
+				if err = stream.Send(&rpc.StreamingOpenResp{
+					Data:    buffer[:bufferUsed],
+					Dropped: int32(dropped),
+				}); err != nil {
+					// error sending to stream
+					streamClosed <- err
+					break
+				}
+				bufferUsed = 0
+				dropped = 0
+
+				// Rate limit, filling all the available window
+				if rateLimitEnabled {
+					slots = atomic.AddInt32(&writeSlots, -1)
+					//fmt.Println("FREE SLOTS:", slots)
+				}
 			}
 		}
 	}()

rpc/monitor/monitor.pb.go

@@ -44,6 +44,10 @@ message StreamingOpenReq {
 
     // The data to be sent to the target being monitored.
     bytes data = 2;
+
+    // When the rate limiter is enabled, this parameter is used to report the
+    // number of successfully processed StreamingOpenResp messages (with data).
+    int32 recv_acknowledge = 3;
   }
 }
 
@@ -61,10 +65,22 @@ message MonitorConfig {
   // Additional parameters that might be needed to configure the target or the
   // monitor itself.
   google.protobuf.Struct additionalConfig = 3;
+
+  // This parameter indicates how many bytes should be buffered on the server side
+  // before dropping. If >0 then the server will enable a rate limiter and will send
+  // incoming data to the client only when the client allows it: see the
+  // StreamingOpenReq.recv_acknowledge parameter for details.
+  int32 recv_rate_limit_buffer = 4;
 }
 
 //
 message StreamingOpenResp {
     // The data received from the target.
     bytes data = 1;
+
+    // The number of bytes dropped.
+    // During regular updates this number should be 0, but in case the
+    // client is not able to process the recv window quickly enough this
+    // parameter will report the number of dropped bytes.
+    int32 dropped = 2;
 }