github.com/ethersphere/bee/v2@v2.2.0/pkg/storer/subscribe_push.go (about)

     1  // Copyright 2023 The Swarm Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package storer
     6  
     7  import (
     8  	"context"
     9  	"sync"
    10  	"time"
    11  
    12  	"github.com/ethersphere/bee/v2/pkg/storer/internal/upload"
    13  	"github.com/ethersphere/bee/v2/pkg/swarm"
    14  )
    15  
    16  const subscribePushEventKey = "subscribe-push"
    17  
    18  func (db *DB) SubscribePush(ctx context.Context) (<-chan swarm.Chunk, func()) {
    19  	chunks := make(chan swarm.Chunk)
    20  
    21  	var (
    22  		stopChan     = make(chan struct{})
    23  		stopChanOnce sync.Once
    24  	)
    25  
    26  	db.subscriptionsWG.Add(1)
    27  	go func() {
    28  		defer db.subscriptionsWG.Done()
    29  
    30  		trigger, unsub := db.events.Subscribe(subscribePushEventKey)
    31  		defer unsub()
    32  
    33  		// close the returned chunkInfo channel at the end to
    34  		// signal that the subscription is done
    35  		defer close(chunks)
    36  		for {
    37  
    38  			err := upload.IteratePending(ctx, db.storage, func(chunk swarm.Chunk) (bool, error) {
    39  				select {
    40  				case chunks <- chunk:
    41  					return false, nil
    42  				case <-stopChan:
    43  					// gracefully stop the iteration
    44  					// on stop
    45  					return true, nil
    46  				case <-db.quit:
    47  					return true, ErrDBQuit
    48  				case <-ctx.Done():
    49  					return true, ctx.Err()
    50  				}
    51  			})
    52  
    53  			if err != nil {
    54  				// if we get storage.ErrNotFound, it could happen that the previous
    55  				// iteration happened on a snapshot that was not fully updated yet.
    56  				// in this case, we wait for the next event to trigger the iteration
    57  				// again. This trigger ensures that we perform the iteration on the
    58  				// latest snapshot.
    59  				db.logger.Error(err, "subscribe push: iterate error")
    60  				select {
    61  				case <-db.quit:
    62  					return
    63  				case <-ctx.Done():
    64  					return
    65  				case <-stopChan:
    66  					return
    67  				case <-time.After(time.Second):
    68  				}
    69  				db.events.Trigger(subscribePushEventKey)
    70  			}
    71  
    72  			select {
    73  			case <-db.quit:
    74  				return
    75  			case <-ctx.Done():
    76  				return
    77  			case <-stopChan:
    78  				return
    79  			case <-trigger:
    80  				// wait for the next event
    81  			}
    82  		}
    83  	}()
    84  
    85  	stop := func() {
    86  		stopChanOnce.Do(func() {
    87  			close(stopChan)
    88  		})
    89  	}
    90  
    91  	return chunks, stop
    92  }