github.com/arieschain/arieschain@v0.0.0-20191023063405-37c074544356/rpc/subscription.go (about)

     1  package rpc
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"sync"
     7  )
     8  
     9  var (
    10  	// ErrNotificationsUnsupported is returned when the connection doesn't support notifications
    11  	ErrNotificationsUnsupported = errors.New("notifications not supported")
    12  	// ErrNotificationNotFound is returned when the notification for the given id is not found
    13  	ErrSubscriptionNotFound = errors.New("subscription not found")
    14  )
    15  
    16  // ID defines a pseudo random number that is used to identify RPC subscriptions.
    17  type ID string
    18  
    19  // a Subscription is created by a notifier and tight to that notifier. The client can use
    20  // this subscription to wait for an unsubscribe request for the client, see Err().
    21  type Subscription struct {
    22  	ID        ID
    23  	namespace string
    24  	err       chan error // closed on unsubscribe
    25  }
    26  
    27  // Err returns a channel that is closed when the client send an unsubscribe request.
    28  func (s *Subscription) Err() <-chan error {
    29  	return s.err
    30  }
    31  
    32  // notifierKey is used to store a notifier within the connection context.
    33  type notifierKey struct{}
    34  
    35  // Notifier is tight to a RPC connection that supports subscriptions.
    36  // Server callbacks use the notifier to send notifications.
    37  type Notifier struct {
    38  	codec    ServerCodec
    39  	subMu    sync.RWMutex // guards active and inactive maps
    40  	active   map[ID]*Subscription
    41  	inactive map[ID]*Subscription
    42  }
    43  
    44  // newNotifier creates a new notifier that can be used to send subscription
    45  // notifications to the client.
    46  func newNotifier(codec ServerCodec) *Notifier {
    47  	return &Notifier{
    48  		codec:    codec,
    49  		active:   make(map[ID]*Subscription),
    50  		inactive: make(map[ID]*Subscription),
    51  	}
    52  }
    53  
    54  // NotifierFromContext returns the Notifier value stored in ctx, if any.
    55  func NotifierFromContext(ctx context.Context) (*Notifier, bool) {
    56  	n, ok := ctx.Value(notifierKey{}).(*Notifier)
    57  	return n, ok
    58  }
    59  
    60  // CreateSubscription returns a new subscription that is coupled to the
    61  // RPC connection. By default subscriptions are inactive and notifications
    62  // are dropped until the subscription is marked as active. This is done
    63  // by the RPC server after the subscription ID is send to the client.
    64  func (n *Notifier) CreateSubscription() *Subscription {
    65  	s := &Subscription{ID: NewID(), err: make(chan error)}
    66  	n.subMu.Lock()
    67  	n.inactive[s.ID] = s
    68  	n.subMu.Unlock()
    69  	return s
    70  }
    71  
    72  // Notify sends a notification to the client with the given data as payload.
    73  // If an error occurs the RPC connection is closed and the error is returned.
    74  func (n *Notifier) Notify(id ID, data interface{}) error {
    75  	n.subMu.RLock()
    76  	defer n.subMu.RUnlock()
    77  
    78  	sub, active := n.active[id]
    79  	if active {
    80  		notification := n.codec.CreateNotification(string(id), sub.namespace, data)
    81  		if err := n.codec.Write(notification); err != nil {
    82  			n.codec.Close()
    83  			return err
    84  		}
    85  	}
    86  	return nil
    87  }
    88  
    89  // Closed returns a channel that is closed when the RPC connection is closed.
    90  func (n *Notifier) Closed() <-chan interface{} {
    91  	return n.codec.Closed()
    92  }
    93  
    94  // unsubscribe a subscription.
    95  // If the subscription could not be found ErrSubscriptionNotFound is returned.
    96  func (n *Notifier) unsubscribe(id ID) error {
    97  	n.subMu.Lock()
    98  	defer n.subMu.Unlock()
    99  	if s, found := n.active[id]; found {
   100  		close(s.err)
   101  		delete(n.active, id)
   102  		return nil
   103  	}
   104  	return ErrSubscriptionNotFound
   105  }
   106  
   107  // activate enables a subscription. Until a subscription is enabled all
   108  // notifications are dropped. This method is called by the RPC server after
   109  // the subscription ID was sent to client. This prevents notifications being
   110  // send to the client before the subscription ID is send to the client.
   111  func (n *Notifier) activate(id ID, namespace string) {
   112  	n.subMu.Lock()
   113  	defer n.subMu.Unlock()
   114  	if sub, found := n.inactive[id]; found {
   115  		sub.namespace = namespace
   116  		n.active[id] = sub
   117  		delete(n.inactive, id)
   118  	}
   119  }