Add a way to specify batch size on response subscription. (#259)

* Add a way to specify batch size on response subscription.

ChangeLog

@@ -1,3 +1,8 @@
+Version 11.3.0
+--------------
+    * Add a way to specify the batch size for the response subscription in
+      Mettle API
+
 Version 11.2.2
 --------------
     * Correctly initialise Prometheus Vector metrics

VERSION

@@ -1 +1 @@
-11.2.2
+11.3.0

mettle/api/api.go

@@ -92,15 +92,15 @@ func init() {
 }
 
 // ServeForever serves on the given port until terminated.
-func ServeForever(opts grpcutil.Opts, name, requestQueue, responseQueue, preResponseQueue, apiURL string, connTLS bool, allowedPlatform map[string][]string, storageURL string, storageTLS bool, numPollers int) {
-	s, lis, err := serve(opts, name, requestQueue, responseQueue, preResponseQueue, apiURL, connTLS, allowedPlatform, storageURL, storageTLS, numPollers)
+func ServeForever(opts grpcutil.Opts, name, requestQueue, responseQueue, preResponseQueue, apiURL string, connTLS bool, allowedPlatform map[string][]string, storageURL string, storageTLS bool, numPollers int, responseBatchSize uint) {
+	s, lis, err := serve(opts, name, requestQueue, responseQueue, preResponseQueue, apiURL, connTLS, allowedPlatform, storageURL, storageTLS, numPollers, responseBatchSize)
 	if err != nil {
 		log.Fatalf("%s", err)
 	}
 	grpcutil.ServeForever(lis, s)
 }
 
-func serve(opts grpcutil.Opts, name, requestQueue, responseQueue, preResponseQueue, apiURL string, connTLS bool, allowedPlatform map[string][]string, storageURL string, storageTLS bool, numPollers int) (*grpc.Server, net.Listener, error) {
+func serve(opts grpcutil.Opts, name, requestQueue, responseQueue, preResponseQueue, apiURL string, connTLS bool, allowedPlatform map[string][]string, storageURL string, storageTLS bool, numPollers int, responseBatchSize uint) (*grpc.Server, net.Listener, error) {
 	if name == "" {
 		name = "mettle API server"
 	}
@@ -115,7 +115,7 @@ func serve(opts grpcutil.Opts, name, requestQueue, responseQueue, preResponseQue
 	srv := &server{
 		name:         name,
 		requests:     common.MustOpenTopic(requestQueue),
-		responses:    common.MustOpenSubscription(responseQueue),
+		responses:    common.MustOpenSubscription(responseQueue, responseBatchSize),
 		preResponses: common.MustOpenTopic(preResponseQueue),
 		jobs:         map[string]*job{},
 		platform:     allowedPlatform,

mettle/api/api_test.go

@@ -245,7 +245,7 @@ func setupServers(t *testing.T) (pb.ExecutionClient, *executor, *grpc.Server) {
 	s, lis, err := serve(grpcutil.Opts{
 		Host: "127.0.0.1",
 		Port: 0,
-	}, "", requests, responses, responses, "", true, map[string][]string{}, casaddr, false, 1)
+	}, "", requests, responses, responses, "", true, map[string][]string{}, casaddr, false, 1, 1)
 	require.NoError(t, err)
 	go s.Serve(lis)
 	conn, err := grpc.Dial(lis.Addr().String(), grpc.WithInsecure())
@@ -328,7 +328,7 @@ type executor struct {
 
 func newExecutor(requests, responses string) *executor {
 	return &executor{
-		requests:  common.MustOpenSubscription(requests),
+		requests:  common.MustOpenSubscription(requests, 1),
 		responses: common.MustOpenTopic(responses),
 	}
 }

mettle/common/common.go

@@ -7,6 +7,7 @@ import (
 	"os"
 	"os/signal"
 	"path/filepath"
+	"strconv"
 	"strings"
 	"sync"
 	"syscall"
@@ -37,9 +38,13 @@ var subMutex sync.Mutex
 const workerKey = "build.please.mettle.worker"
 
 // MustOpenSubscription opens a subscription, which must have been created ahead of time.
-// It dies on any errors.
-func MustOpenSubscription(url string) *pubsub.Subscription {
-	url = limitBatchSize(url, "1")
+// It dies on any errors. `batchSize` controls how many messages can receive at a time.
+// It can not be equal to zero; this function will panic if it is.
+func MustOpenSubscription(url string, batchSize uint) *pubsub.Subscription {
+	if batchSize == 0 {
+		log.Fatal("batch size for subscription can not be equal to zero")
+	}
+	url = limitBatchSize(url, strconv.Itoa(int(batchSize)))
 	subMutex.Lock()
 	defer subMutex.Unlock()
 	if sub, present := subscriptions[url]; present {

mettle/main.go

@@ -60,6 +60,7 @@ var opts = struct {
 			ResponseQueueSuffix string `long:"response_queue_suffix" env:"API_RESPONSE_QUEUE_SUFFIX" description:"Suffix to apply to the response queue name"`
 			PreResponseQueue    string `long:"pre_response_queue" env:"API_PRE_RESPONSE_QUEUE" required:"true" description:"URL describing the pub/sub queue to connect to for preloading responses to other servers"`
 			NumPollers          int    `long:"num_pollers" env:"API_NUM_POLLERS" default:"10"`
+			ResponseBatchSize   uint   `long:"response_batch_size" env:"API_RESPONSE_BATCH_SIZE" default:"100"`
 		} `group:"Options controlling the pub/sub queues"`
 		AllowedPlatform map[string][]string `long:"allowed_platform" description:"Allowed values for platform properties"`
 	} `command:"api" description:"Start as an API server"`
@@ -179,11 +180,11 @@ func main() {
 			storage := opts.Dual.Storage.Storage[i%len(opts.Dual.Storage.Storage)]
 			go worker.RunForever(opts.InstanceName, requests+"?ackdeadline=10m", responses, fmt.Sprintf("%s-%d", opts.InstanceName, i), storage, opts.Dual.Dir, opts.Dual.Cache.Dir, opts.Dual.Browser, opts.Dual.Sandbox, opts.Dual.AltSandbox, opts.Dual.Lucidity, "", opts.Dual.GRPC.TokenFile, opts.Dual.Redis.URL, opts.Dual.Redis.ReadURL, opts.Dual.Redis.ReadPassword(), opts.Dual.Redis.CAFile, opts.Dual.Redis.TLS, opts.Dual.Cache.Prefix, opts.Dual.Cache.Part, !opts.Dual.NoClean, opts.Dual.Storage.TLS, int64(opts.Dual.Cache.MaxMem), int64(opts.Dual.MinDiskSpace), opts.Dual.MemoryThreshold, opts.Dual.VersionFile, opts.Dual.Costs, 0, opts.Worker.ImmediateShutdown)
 		}
-		api.ServeForever(opts.Dual.GRPC, "", requests, responses, responses, "", false, opts.Dual.AllowedPlatform, opts.Dual.Storage.Storage[0], opts.Dual.Storage.TLS, opts.API.Queues.NumPollers)
+		api.ServeForever(opts.Dual.GRPC, "", requests, responses, responses, "", false, opts.Dual.AllowedPlatform, opts.Dual.Storage.Storage[0], opts.Dual.Storage.TLS, opts.API.Queues.NumPollers, opts.API.Queues.ResponseBatchSize)
 	} else if cmd == "worker" {
 		worker.RunForever(opts.InstanceName, opts.Worker.Queues.RequestQueue, opts.Worker.Queues.ResponseQueue, opts.Worker.Name, opts.Worker.Storage.Storage, opts.Worker.Dir, opts.Worker.Cache.Dir, opts.Worker.Browser, opts.Worker.Sandbox, opts.Worker.AltSandbox, opts.Worker.Lucidity, opts.Worker.PromGateway, opts.Worker.Storage.TokenFile, opts.Worker.Redis.URL, opts.Worker.Redis.ReadURL, opts.Worker.Redis.ReadPassword(), opts.Worker.Redis.CAFile, opts.Worker.Redis.TLS, opts.Worker.Cache.Prefix, opts.Worker.Cache.Part, !opts.Worker.NoClean, opts.Worker.Storage.TLS, int64(opts.Worker.Cache.MaxMem), int64(opts.Worker.MinDiskSpace), opts.Worker.MemoryThreshold, opts.Worker.VersionFile, opts.Worker.Costs, time.Duration(opts.Worker.Queues.AckExtension), opts.Worker.ImmediateShutdown)
 	} else if cmd == "api" {
-		api.ServeForever(opts.API.GRPC, opts.API.Queues.ResponseQueueSuffix, opts.API.Queues.RequestQueue, opts.API.Queues.ResponseQueue+opts.API.Queues.ResponseQueueSuffix, opts.API.Queues.PreResponseQueue, opts.API.API.URL, opts.API.API.TLS, opts.API.AllowedPlatform, opts.API.Storage.Storage, opts.API.Storage.TLS, opts.API.Queues.NumPollers)
+		api.ServeForever(opts.API.GRPC, opts.API.Queues.ResponseQueueSuffix, opts.API.Queues.RequestQueue, opts.API.Queues.ResponseQueue+opts.API.Queues.ResponseQueueSuffix, opts.API.Queues.PreResponseQueue, opts.API.API.URL, opts.API.API.TLS, opts.API.AllowedPlatform, opts.API.Storage.Storage, opts.API.Storage.TLS, opts.API.Queues.NumPollers, opts.API.Queues.ResponseBatchSize)
 	} else if err := one(); err != nil {
 		log.Fatalf("%s", err)
 	}

mettle/worker/worker.go

@@ -289,7 +289,7 @@ func initialiseWorker(instanceName, requestQueue, responseQueue, name, storage,
 	}
 
 	w := &worker{
-		requests:        common.MustOpenSubscription(requestQueue),
+		requests:        common.MustOpenSubscription(requestQueue, 1),
 		responses:       common.MustOpenTopic(responseQueue),
 		ackExtension:    ackExtension,
 		client:          client,