Add support for bytes limit

cmd/msak-client/client.go

@@ -15,15 +15,16 @@ const clientName = "msak-client-go"
 var clientVersion = version.Version
 
 var (
-	flagServer   = flag.String("server", "", "Server address")
-	flagStreams  = flag.Int("streams", client.DefaultStreams, "Number of streams")
-	flagCC       = flag.String("cc", "bbr", "Congestion control algorithm to use")
-	flagDelay    = flag.Duration("delay", 0, "Delay between each stream")
-	flagDuration = flag.Duration("duration", client.DefaultLength, "Length of the last stream")
-	flagScheme   = flag.String("scheme", client.DefaultScheme, "Websocket scheme (wss or ws)")
-	flagMID      = flag.String("mid", uuid.NewString(), "Measurement ID to use")
-	flagNoVerify = flag.Bool("no-verify", false, "Skip TLS certificate verification")
-	flagDebug    = flag.Bool("debug", false, "Enable debug logging")
+	flagServer     = flag.String("server", "", "Server address")
+	flagStreams    = flag.Int("streams", client.DefaultStreams, "Number of streams")
+	flagCC         = flag.String("cc", "bbr", "Congestion control algorithm to use")
+	flagDelay      = flag.Duration("delay", 0, "Delay between each stream")
+	flagDuration   = flag.Duration("duration", client.DefaultLength, "Length of the last stream")
+	flagScheme     = flag.String("scheme", client.DefaultScheme, "Websocket scheme (wss or ws)")
+	flagMID        = flag.String("mid", uuid.NewString(), "Measurement ID to use")
+	flagNoVerify   = flag.Bool("no-verify", false, "Skip TLS certificate verification")
+	flagDebug      = flag.Bool("debug", false, "Enable debug logging")
+	flagByteLimit = flag.Int("bytes", 0, "Byte limit to request to the server")
 )
 
 func main() {
@@ -46,7 +47,8 @@ func main() {
 		Emitter: client.HumanReadable{
 			Debug: *flagDebug,
 		},
-		NoVerify: *flagNoVerify,
+		NoVerify:   *flagNoVerify,
+		ByteLimit: *flagByteLimit,
 	}
 
 	cl := client.New(clientName, clientVersion, config)

internal/handler/handler.go

@@ -16,6 +16,7 @@ import (
 	"github.com/m-lab/msak/internal/persistence"
 	"github.com/m-lab/msak/pkg/throughput1"
 	"github.com/m-lab/msak/pkg/throughput1/model"
+	"github.com/m-lab/msak/pkg/throughput1/spec"
 	"github.com/m-lab/msak/pkg/version"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promauto"
@@ -130,6 +131,20 @@ func (h *Handler) upgradeAndRunMeasurement(kind model.TestDirection, rw http.Res
 			model.NameValue{Name: "delay", Value: requestDelay})
 	}
 
+	requestByteLimit := query.Get(spec.ByteLimitParameterName)
+	var byteLimit int
+	if requestByteLimit != "" {
+		if byteLimit, err = strconv.Atoi(requestByteLimit); err != nil {
+			ClientConnections.WithLabelValues(string(kind), "invalid-byte-limit").Inc()
+			log.Info("Received request with an invalid byte limit", "source", req.RemoteAddr,
+				"value", requestByteLimit)
+			writeBadRequest(rw)
+			return
+		}
+		clientOptions = append(clientOptions,
+			model.NameValue{Name: spec.ByteLimitParameterName, Value: requestByteLimit})
+	}
+
 	// Read metadata (i.e. everything in the querystring that's not a known
 	// option).
 	metadata, err := getRequestMetadata(req)
@@ -198,6 +213,7 @@ func (h *Handler) upgradeAndRunMeasurement(kind model.TestDirection, rw http.Res
 	defer cancel()
 
 	proto := throughput1.New(wsConn)
+	proto.SetByteLimit(byteLimit)
 	var senderCh, receiverCh <-chan model.WireMeasurement
 	var errCh <-chan error
 	if kind == model.DirectionDownload {

internal/handler/handler_test.go

@@ -215,6 +215,11 @@ func TestHandler_Validation(t *testing.T) {
 			target:     "/?mid=test&streams=2&duration=invalid",
 			statusCode: http.StatusBadRequest,
 		},
+		{
+			name:       "invalid byte limit",
+			target:     "/?mid=test&streams=2&duration=1000&bytes=invalid",
+			statusCode: http.StatusBadRequest,
+		},
 		{
 			name:       "metadata key too long",
 			target:     "/?mid=test&streams=2&" + longKey,

pkg/client/client.go

@@ -136,6 +136,7 @@ func (c *Throughput1Client) connect(ctx context.Context, serviceURL *url.URL) (*
 	q := serviceURL.Query()
 	q.Set("streams", fmt.Sprint(c.config.NumStreams))
 	q.Set("cc", c.config.CongestionControl)
+	q.Set(spec.ByteLimitParameterName, fmt.Sprint(c.config.ByteLimit))
 	q.Set("duration", fmt.Sprintf("%d", c.config.Length.Milliseconds()))
 	q.Set("client_arch", runtime.GOARCH)
 	q.Set("client_library_name", libraryName)

pkg/client/config.go

@@ -35,4 +35,8 @@ type Config struct {
 
 	// NoVerify disables the TLS certificate verification.
 	NoVerify bool
+
+	// ByteLimit is the maximum number of bytes to download or upload. If set to 0, the
+	// limit is disabled.
+	ByteLimit int
 }

pkg/throughput1/protocol.go

@@ -47,6 +47,8 @@ type Protocol struct {
 
 	applicationBytesReceived atomic.Int64
 	applicationBytesSent     atomic.Int64
+
+	byteLimit int64
 }
 
 // New returns a new Protocol with the specified connection and every other
@@ -61,6 +63,12 @@ func New(conn *websocket.Conn) *Protocol {
 	}
 }
 
+// SetByteLimit sets the number of bytes sent after which a test (either download or upload) will stop.
+// Set the value to zero to disable the byte limit.
+func (p *Protocol) SetByteLimit(value int) {
+	p.byteLimit = int64(value)
+}
+
 // Upgrade takes a HTTP request and upgrades the connection to WebSocket.
 // Returns a websocket Conn if the upgrade succeeded, and an error otherwise.
 func Upgrade(w http.ResponseWriter, r *http.Request) (*websocket.Conn, error) {
@@ -218,6 +226,12 @@ func (p *Protocol) sendCounterflow(ctx context.Context,
 			case results <- wm:
 			default:
 			}
+
+			// End the test once enough bytes have been received.
+			if p.byteLimit > 0 && m.TCPInfo != nil && m.TCPInfo.BytesReceived >= p.byteLimit {
+				p.close(ctx)
+				return
+			}
 		}
 	}
 }
@@ -274,6 +288,12 @@ func (p *Protocol) sender(ctx context.Context, measurerCh <-chan model.Measureme
 			case results <- wm:
 			default:
 			}
+
+			// End the test once enough bytes have been acked.
+			if p.byteLimit > 0 && m.TCPInfo != nil && m.TCPInfo.BytesAcked >= p.byteLimit {
+				p.close(ctx)
+				return
+			}
 		default:
 			err = p.conn.WritePreparedMessage(message)
 			if err != nil {

pkg/throughput1/spec/spec.go

@@ -37,6 +37,11 @@ const (
 
 	// SecWebSocketProtocol is the value of the Sec-WebSocket-Protocol header.
 	SecWebSocketProtocol = "net.measurementlab.throughput.v1"
+
+	// ByteLimitParameterName is the name of the parameter that clients can use
+	// to terminate throughput1 download tests once the test has transferred
+	// the specified number of bytes.
+	ByteLimitParameterName = "bytes"
 )
 
 // SubtestKind indicates the subtest kind