From a261c2ffe18220cce1bd290e23e67405ba704db9 Mon Sep 17 00:00:00 2001
From: AndrewSisley <a.j.sisley@hotmail.com>
Date: Thu, 30 Mar 2023 21:36:43 -0400
Subject: [PATCH] test: Fix flaky TestCloseThroughContext test (#1265)

---
 datastore/memory/memory.go      | 12 +++++-------
 datastore/memory/memory_test.go |  2 +-
 2 files changed, 6 insertions(+), 8 deletions(-)

diff --git a/datastore/memory/memory.go b/datastore/memory/memory.go
index 5080b157ee..277c24a0fc 100644
--- a/datastore/memory/memory.go
+++ b/datastore/memory/memory.go
@@ -70,10 +70,9 @@ type Datastore struct {
 	inFlightTxn *btree.BTreeG[dsTxn]
 	commit      chan commit
 
-	closing   chan struct{}
-	closed    bool
-	closeLk   sync.RWMutex
-	closeOnce sync.Once
+	closing chan struct{}
+	closed  bool
+	closeLk sync.RWMutex
 }
 
 var _ ds.Datastore = (*Datastore)(nil)
@@ -117,15 +116,14 @@ func (d *Datastore) newBasicBatch() ds.Batch {
 }
 
 func (d *Datastore) Close() error {
-	d.closeOnce.Do(func() {
-		close(d.closing)
-	})
 	d.closeLk.Lock()
 	defer d.closeLk.Unlock()
 	if d.closed {
 		return ErrClosed
 	}
+
 	d.closed = true
+	close(d.closing)
 	close(d.commit)
 	return nil
 }
diff --git a/datastore/memory/memory_test.go b/datastore/memory/memory_test.go
index 02cdb3b598..1dbf7a825e 100644
--- a/datastore/memory/memory_test.go
+++ b/datastore/memory/memory_test.go
@@ -92,7 +92,7 @@ func TestCloseThroughContext(t *testing.T) {
 
 	cancel()
 
-	<-s.closing
+	time.Sleep(time.Millisecond * 10)
 
 	err := s.Close()
 	require.ErrorIs(t, err, ErrClosed)