Refactor client info manager (Velocidex#1700)

* Refactor client info manager

* Uses larger mutations in order to minimize number of updates.
* Flushes client stats less frequently to minimize IO

* Added test

api/api.go

@@ -66,10 +66,10 @@ import (
 
 type ApiServer struct {
 	proto.UnimplementedAPIServer
-	config     *config_proto.Config
-	server_obj *server.Server
-	ca_pool    *x509.CertPool
-
+	config             *config_proto.Config
+	server_obj         *server.Server
+	ca_pool            *x509.CertPool
+	wg                 *sync.WaitGroup
 	api_client_factory grpc_client.APIClientFactory
 }
 
@@ -984,6 +984,7 @@ func startAPIServer(
 			server_obj:         server_obj,
 			ca_pool:            CA_Pool,
 			api_client_factory: grpc_client.GRPCAPIClient{},
+			wg:                 wg,
 		},
 	)
 	// Register reflection service.

api/proto/api.proto

@@ -75,6 +75,8 @@ message EventRequest {
 
     // The node who is requesting the event stream.
     string node = 2;
+
+    string watcher_name = 3;
 }
 
 message EventResponse {

api/replication.go

@@ -63,28 +63,35 @@ func streamEvents(
 
 	// The API service is running on the master only! This means
 	// the journal service is local.
-	output_chan, cancel := journal.Watch(ctx, in.Queue)
+	output_chan, cancel := journal.Watch(
+		ctx, in.Queue, "replication-"+in.WatcherName)
 	defer cancel()
 
-	for event := range output_chan {
-		serialized, err := json.Marshal(event)
-		if err != nil {
-			continue
-		}
-		response := &api_proto.EventResponse{
-			Jsonl: serialized,
-		}
-
-		timer := prometheus.NewTimer(
-			prometheus.ObserverFunc(func(v float64) {
-				replicationReceiveHistorgram.WithLabelValues("").Observe(v)
-			}))
-
-		err = stream.Send(response)
-		timer.ObserveDuration()
-
-		if err != nil {
-			continue
+	for {
+		select {
+		case <-ctx.Done():
+			return
+
+		case event := <-output_chan:
+			serialized, err := json.Marshal(event)
+			if err != nil {
+				continue
+			}
+			response := &api_proto.EventResponse{
+				Jsonl: serialized,
+			}
+
+			timer := prometheus.NewTimer(
+				prometheus.ObserverFunc(func(v float64) {
+					replicationReceiveHistorgram.WithLabelValues("").Observe(v)
+				}))
+
+			err = stream.Send(response)
+			timer.ObserveDuration()
+
+			if err != nil {
+				continue
+			}
 		}
 	}
 
@@ -140,8 +147,13 @@ func (self *ApiServer) WatchEvent(
 
 		// return the first good match
 		if true {
+			// Wait here for orderly shutdown of event streams.
+			self.wg.Add(1)
+			defer self.wg.Done()
+
 			// Cert is good enough for us, run the query.
-			return streamEvents(ctx, self.config, in, stream, peer_name)
+			return streamEvents(
+				ctx, self.config, in, stream, peer_name)
 		}
 	}
 

artifacts/definitions/Server/Import/PreviousReleases.yaml

@@ -23,6 +23,7 @@ parameters:
     description: |
       The Velociraptor Release to import.
     type: choices
+    default: v0.6.3
     choices:
       - v0.6.0
       - v0.6.1