fix: close iterator

This commit is contained in:
Richard Ramos 2023-01-31 10:45:19 -04:00 committed by RichΛrd
parent 5ff87d8e65
commit 6169a44c24
2 changed files with 29 additions and 2 deletions

View File

@ -256,7 +256,20 @@ func (d *DiscoveryV5) iterate(ctx context.Context) error {
return fmt.Errorf("obtaining iterator: %w", err)
}
defer iterator.Close()
closeCh := make(chan struct{}, 1)
defer close(closeCh)
// Closing iterator when context is cancelled or function is returning
d.wg.Add(1)
go func() {
defer d.wg.Done()
select {
case <-ctx.Done():
iterator.Close()
case <-closeCh:
iterator.Close()
}
}()
for {
if ctx.Err() != nil {

View File

@ -317,7 +317,21 @@ func (wakuPX *WakuPeerExchange) iterate(ctx context.Context) error {
if err != nil {
return fmt.Errorf("obtaining iterator: %w", err)
}
defer iterator.Close()
closeCh := make(chan struct{}, 1)
defer close(closeCh)
// Closing iterator when context is cancelled or function is returning
wakuPX.wg.Add(1)
go func() {
defer wakuPX.wg.Done()
select {
case <-ctx.Done():
iterator.Close()
case <-closeCh:
iterator.Close()
}
}()
for {
if ctx.Err() != nil {