github.com/keltia/go-ipfs@v0.3.8-0.20150909044612-210793031c63/exchange/bitswap/notifications/notifications.go (about)

     1  package notifications
     2  
     3  import (
     4  	pubsub "github.com/ipfs/go-ipfs/Godeps/_workspace/src/github.com/briantigerchow/pubsub"
     5  	context "github.com/ipfs/go-ipfs/Godeps/_workspace/src/golang.org/x/net/context"
     6  	blocks "github.com/ipfs/go-ipfs/blocks"
     7  	key "github.com/ipfs/go-ipfs/blocks/key"
     8  )
     9  
    10  const bufferSize = 16
    11  
    12  type PubSub interface {
    13  	Publish(block *blocks.Block)
    14  	Subscribe(ctx context.Context, keys ...key.Key) <-chan *blocks.Block
    15  	Shutdown()
    16  }
    17  
    18  func New() PubSub {
    19  	return &impl{*pubsub.New(bufferSize)}
    20  }
    21  
    22  type impl struct {
    23  	wrapped pubsub.PubSub
    24  }
    25  
    26  func (ps *impl) Publish(block *blocks.Block) {
    27  	topic := string(block.Key())
    28  	ps.wrapped.Pub(block, topic)
    29  }
    30  
    31  func (ps *impl) Shutdown() {
    32  	ps.wrapped.Shutdown()
    33  }
    34  
    35  // Subscribe returns a channel of blocks for the given |keys|. |blockChannel|
    36  // is closed if the |ctx| times out or is cancelled, or after sending len(keys)
    37  // blocks.
    38  func (ps *impl) Subscribe(ctx context.Context, keys ...key.Key) <-chan *blocks.Block {
    39  
    40  	blocksCh := make(chan *blocks.Block, len(keys))
    41  	valuesCh := make(chan interface{}, len(keys)) // provide our own channel to control buffer, prevent blocking
    42  	if len(keys) == 0 {
    43  		close(blocksCh)
    44  		return blocksCh
    45  	}
    46  	ps.wrapped.AddSubOnceEach(valuesCh, toStrings(keys)...)
    47  	go func() {
    48  		defer close(blocksCh)
    49  		defer ps.wrapped.Unsub(valuesCh) // with a len(keys) buffer, this is an optimization
    50  		for {
    51  			select {
    52  			case <-ctx.Done():
    53  				return
    54  			case val, ok := <-valuesCh:
    55  				if !ok {
    56  					return
    57  				}
    58  				block, ok := val.(*blocks.Block)
    59  				if !ok {
    60  					return
    61  				}
    62  				select {
    63  				case <-ctx.Done():
    64  					return
    65  				case blocksCh <- block: // continue
    66  				}
    67  			}
    68  		}
    69  	}()
    70  
    71  	return blocksCh
    72  }
    73  
    74  func toStrings(keys []key.Key) []string {
    75  	strs := make([]string, 0)
    76  	for _, key := range keys {
    77  		strs = append(strs, string(key))
    78  	}
    79  	return strs
    80  }