From ddd472719b8d7bc8c9ff3195f03d84a615bc0392 Mon Sep 17 00:00:00 2001
From: Fae Charlton <fae.charlton@elastic.co>
Date: Tue, 21 Jul 2020 15:32:16 -0400
Subject: [PATCH 1/2] Fix write error in EnsureWriter, add test to confirm

---
 libbeat/statestore/backend/memlog/util.go     |  2 +-
 .../statestore/backend/memlog/util_test.go    | 47 +++++++++++++++++++
 2 files changed, 48 insertions(+), 1 deletion(-)
 create mode 100644 libbeat/statestore/backend/memlog/util_test.go

diff --git a/libbeat/statestore/backend/memlog/util.go b/libbeat/statestore/backend/memlog/util.go
index 2027c87adca..e2c5d4e6f68 100644
--- a/libbeat/statestore/backend/memlog/util.go
+++ b/libbeat/statestore/backend/memlog/util.go
@@ -55,7 +55,7 @@ func (e *ensureWriter) Write(p []byte) (int, error) {
 	for len(p) > 0 {
 		n, err := e.w.Write(p)
 		N, p = N+n, p[n:]
-		if isRetryErr(err) {
+		if err != nil && !isRetryErr(err) {
 			return N, err
 		}
 	}
diff --git a/libbeat/statestore/backend/memlog/util_test.go b/libbeat/statestore/backend/memlog/util_test.go
new file mode 100644
index 00000000000..b960f33d5a9
--- /dev/null
+++ b/libbeat/statestore/backend/memlog/util_test.go
@@ -0,0 +1,47 @@
+// Licensed to Elasticsearch B.V. under one or more contributor
+// license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright
+// ownership. Elasticsearch B.V. licenses this file to you under
+// the Apache License, Version 2.0 (the "License"); you may
+// not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package memlog
+
+import (
+	"syscall"
+	"testing"
+)
+
+type mockErrorWriter struct {
+	reportedError bool
+}
+
+func (mew *mockErrorWriter) Write(data []byte) (n int, err error) {
+	if !mew.reportedError {
+		mew.reportedError = true
+		return 0, syscall.EAGAIN
+	}
+	return len(data), nil
+}
+
+func TestEnsureWriter_RetriableError(t *testing.T) {
+	errorWriter := &mockErrorWriter{}
+	writer := &ensureWriter{errorWriter}
+	written, err := writer.Write([]byte{1, 2, 3})
+	if err != nil {
+		t.Fatalf("EnsureWriter shouldn't propagate retriable errors")
+	}
+	if written != 3 {
+		t.Fatalf("Expected 3 bytes written, got %d", written)
+	}
+}

From dde8b9548db3981066bc78eb8ce49cc33a8c55a2 Mon Sep 17 00:00:00 2001
From: Fae Charlton <fae.charlton@elastic.co>
Date: Tue, 21 Jul 2020 15:53:31 -0400
Subject: [PATCH 2/2] Add more test cases

---
 .../statestore/backend/memlog/util_test.go    | 46 ++++++++++++++++---
 1 file changed, 40 insertions(+), 6 deletions(-)

diff --git a/libbeat/statestore/backend/memlog/util_test.go b/libbeat/statestore/backend/memlog/util_test.go
index b960f33d5a9..fca2a2bbaf6 100644
--- a/libbeat/statestore/backend/memlog/util_test.go
+++ b/libbeat/statestore/backend/memlog/util_test.go
@@ -22,26 +22,60 @@ import (
 	"testing"
 )
 
+// A mock Writer implementation that always returns a configurable
+// error on the first write call, to test error handling in ensureWriter.
 type mockErrorWriter struct {
+	errorType     error
 	reportedError bool
 }
 
 func (mew *mockErrorWriter) Write(data []byte) (n int, err error) {
 	if !mew.reportedError {
 		mew.reportedError = true
-		return 0, syscall.EAGAIN
+		return 0, mew.errorType
 	}
 	return len(data), nil
 }
 
 func TestEnsureWriter_RetriableError(t *testing.T) {
-	errorWriter := &mockErrorWriter{}
+	// EAGAIN is retriable, ensureWriter.Write should succeed.
+	errorWriter := &mockErrorWriter{errorType: syscall.EAGAIN}
+	bytes := []byte{1, 2, 3}
 	writer := &ensureWriter{errorWriter}
-	written, err := writer.Write([]byte{1, 2, 3})
+	written, err := writer.Write(bytes)
 	if err != nil {
-		t.Fatalf("EnsureWriter shouldn't propagate retriable errors")
+		t.Fatalf("ensureWriter shouldn't propagate retriable errors")
 	}
-	if written != 3 {
-		t.Fatalf("Expected 3 bytes written, got %d", written)
+	if written != len(bytes) {
+		t.Fatalf("Expected %d bytes written, got %d", len(bytes), written)
+	}
+}
+
+func TestEnsureWriter_NonRetriableError(t *testing.T) {
+	// EINVAL is not retriable, ensureWriter.Write should return an error.
+	errorWriter := &mockErrorWriter{errorType: syscall.EINVAL}
+	bytes := []byte{1, 2, 3}
+	writer := &ensureWriter{errorWriter}
+	written, err := writer.Write(bytes)
+	if err != syscall.EINVAL {
+		t.Fatalf("ensureWriter should propagate nonretriable errors")
+	}
+	if written != 0 {
+		t.Fatalf("Expected 0 bytes written, got %d", written)
+	}
+}
+
+func TestEnsureWriter_NoError(t *testing.T) {
+	// This tests the case where the underlying writer returns with no error,
+	// but without writing the full buffer.
+	var bytes []byte = []byte{1, 2, 3}
+	errorWriter := &mockErrorWriter{errorType: nil}
+	writer := &ensureWriter{errorWriter}
+	written, err := writer.Write(bytes)
+	if err != nil {
+		t.Fatalf("ensureWriter should only error if the underlying writer does")
+	}
+	if written != len(bytes) {
+		t.Fatalf("Expected %d bytes written, got %d", len(bytes), written)
 	}
 }