test: Fix flaky TestCloseThroughContext test (sourcenetwork#1265)

shahzadlone · Mar 31, 2023 · fd599d4 · fd599d4
1 parent b37a544
commit fd599d4
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 8 deletions.
diff --git 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
@@ -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)