Avoid returning incomplete stats from dispatcher

3 files changed, 27 insertions, 4 deletions

diff --git a/client/go/internal/vespa/document/dispatcher.go b/client/go/internal/vespa/document/dispatcher.go
index ca8c585a295..a2b84aaeef2 100644
--- a/client/go/internal/vespa/document/dispatcher.go
+++ b/client/go/internal/vespa/document/dispatcher.go
@@ -157,8 +157,7 @@ func (d *Dispatcher) dispatchNext(id Id) {
 	}
 	hasNext := q != nil
 	if hasNext {
-		next, ok := q.Poll()
-		if ok {
+		if next, ok := q.Poll(); ok {
 			// we have more operations with this ID: dispatch the next one
 			d.dispatch(next)
 		} else {
@@ -240,8 +239,9 @@ func (d *Dispatcher) Enqueue(doc Document) error { return d.enqueue(documentOp{d
 func (d *Dispatcher) Stats() Stats {
 	d.statsMu.Lock()
 	defer d.statsMu.Unlock()
-	d.stats.Inflight = d.inflightCount.Load()
-	return d.stats
+	statsCopy := d.stats.Clone()
+	statsCopy.Inflight = d.inflightCount.Load()
+	return statsCopy
 }
 
 // Close waits for all inflight operations to complete and closes the dispatcher.
diff --git a/client/go/internal/vespa/document/stats.go b/client/go/internal/vespa/document/stats.go
index 4865648a826..7696648f703 100644
--- a/client/go/internal/vespa/document/stats.go
+++ b/client/go/internal/vespa/document/stats.go
@@ -69,6 +69,17 @@ func (s Stats) Successful() int64 {
 
 func (s Stats) Unsuccessful() int64 { return s.Requests - s.Successful() }
 
+func (s Stats) Clone() Stats {
+	if s.ResponsesByCode != nil {
+		mapCopy := make(map[int]int64)
+		for k, v := range s.ResponsesByCode {
+			mapCopy[k] = v
+		}
+		s.ResponsesByCode = mapCopy
+	}
+	return s
+}
+
 // Add statistics from result to this.
 func (s *Stats) Add(result Result) {
 	s.Requests++
diff --git a/client/go/internal/vespa/document/stats_test.go b/client/go/internal/vespa/document/stats_test.go
index 3999ef2e503..8788836f9ad 100644
--- a/client/go/internal/vespa/document/stats_test.go
+++ b/client/go/internal/vespa/document/stats_test.go
@@ -29,3 +29,15 @@ func TestStatsAdd(t *testing.T) {
 		t.Errorf("got stats.Unsuccessful() = %d, want %d", got, want)
 	}
 }
+
+func TestStatsClone(t *testing.T) {
+	var a Stats
+	a.Add(Result{HTTPStatus: 200})
+	b := a.Clone()
+	a.Add(Result{HTTPStatus: 200})
+
+	want := Stats{Requests: 1, Responses: 1, ResponsesByCode: map[int]int64{200: 1}}
+	if !reflect.DeepEqual(b, want) {
+		t.Errorf("got %+v, want %+v", b, want)
+	}
+}