github.com/vantum/vantum@v0.0.0-20180815184342-fe37d5f7a990/rpc/subscription.go (about)

     1  // Copyright 2016 The go-ethereum Authors
     2  // This file is part of the go-ethereum library.
     3  //
     4  // The go-ethereum library is free software: you can redistribute it and/or modify
     5  // it under the terms of the GNU Lesser General Public License as published by
     6  // the Free Software Foundation, either version 3 of the License, or
     7  // (at your option) any later version.
     8  //
     9  // The go-ethereum library is distributed in the hope that it will be useful,
    10  // but WITHOUT ANY WARRANTY; without even the implied warranty of
    11  // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
    12  // GNU Lesser General Public License for more details.
    13  //
    14  // You should have received a copy of the GNU Lesser General Public License
    15  // along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.
    16  
    17  package rpc
    18  
    19  import (
    20  	"context"
    21  	"errors"
    22  	"sync"
    23  )
    24  
    25  var (
    26  	// ErrNotificationsUnsupported is returned when the connection doesn't support notifications
    27  	ErrNotificationsUnsupported = errors.New("notifications not supported")
    28  	// ErrNotificationNotFound is returned when the notification for the given id is not found
    29  	ErrSubscriptionNotFound = errors.New("subscription not found")
    30  )
    31  
    32  // ID defines a pseudo random number that is used to identify RPC subscriptions.
    33  type ID string
    34  
    35  // a Subscription is created by a notifier and tight to that notifier. The client can use
    36  // this subscription to wait for an unsubscribe request for the client, see Err().
    37  type Subscription struct {
    38  	ID        ID
    39  	namespace string
    40  	err       chan error // closed on unsubscribe
    41  }
    42  
    43  // Err returns a channel that is closed when the client send an unsubscribe request.
    44  func (s *Subscription) Err() <-chan error {
    45  	return s.err
    46  }
    47  
    48  // notifierKey is used to store a notifier within the connection context.
    49  type notifierKey struct{}
    50  
    51  // Notifier is tight to a RPC connection that supports subscriptions.
    52  // Server callbacks use the notifier to send notifications.
    53  type Notifier struct {
    54  	codec    ServerCodec
    55  	subMu    sync.RWMutex // guards active and inactive maps
    56  	active   map[ID]*Subscription
    57  	inactive map[ID]*Subscription
    58  }
    59  
    60  // newNotifier creates a new notifier that can be used to send subscription
    61  // notifications to the client.
    62  func newNotifier(codec ServerCodec) *Notifier {
    63  	return &Notifier{
    64  		codec:    codec,
    65  		active:   make(map[ID]*Subscription),
    66  		inactive: make(map[ID]*Subscription),
    67  	}
    68  }
    69  
    70  // NotifierFromContext returns the Notifier value stored in ctx, if any.
    71  func NotifierFromContext(ctx context.Context) (*Notifier, bool) {
    72  	n, ok := ctx.Value(notifierKey{}).(*Notifier)
    73  	return n, ok
    74  }
    75  
    76  // CreateSubscription returns a new subscription that is coupled to the
    77  // RPC connection. By default subscriptions are inactive and notifications
    78  // are dropped until the subscription is marked as active. This is done
    79  // by the RPC server after the subscription ID is send to the client.
    80  func (n *Notifier) CreateSubscription() *Subscription {
    81  	s := &Subscription{ID: NewID(), err: make(chan error)}
    82  	n.subMu.Lock()
    83  	n.inactive[s.ID] = s
    84  	n.subMu.Unlock()
    85  	return s
    86  }
    87  
    88  // Notify sends a notification to the client with the given data as payload.
    89  // If an error occurs the RPC connection is closed and the error is returned.
    90  func (n *Notifier) Notify(id ID, data interface{}) error {
    91  	n.subMu.RLock()
    92  	defer n.subMu.RUnlock()
    93  
    94  	sub, active := n.active[id]
    95  	if active {
    96  		notification := n.codec.CreateNotification(string(id), sub.namespace, data)
    97  		if err := n.codec.Write(notification); err != nil {
    98  			n.codec.Close()
    99  			return err
   100  		}
   101  	}
   102  	return nil
   103  }
   104  
   105  // Closed returns a channel that is closed when the RPC connection is closed.
   106  func (n *Notifier) Closed() <-chan interface{} {
   107  	return n.codec.Closed()
   108  }
   109  
   110  // unsubscribe a subscription.
   111  // If the subscription could not be found ErrSubscriptionNotFound is returned.
   112  func (n *Notifier) unsubscribe(id ID) error {
   113  	n.subMu.Lock()
   114  	defer n.subMu.Unlock()
   115  	if s, found := n.active[id]; found {
   116  		close(s.err)
   117  		delete(n.active, id)
   118  		return nil
   119  	}
   120  	return ErrSubscriptionNotFound
   121  }
   122  
   123  // activate enables a subscription. Until a subscription is enabled all
   124  // notifications are dropped. This method is called by the RPC server after
   125  // the subscription ID was sent to client. This prevents notifications being
   126  // send to the client before the subscription ID is send to the client.
   127  func (n *Notifier) activate(id ID, namespace string) {
   128  	n.subMu.Lock()
   129  	defer n.subMu.Unlock()
   130  	if sub, found := n.inactive[id]; found {
   131  		sub.namespace = namespace
   132  		n.active[id] = sub
   133  		delete(n.inactive, id)
   134  	}
   135  }