github.com/ylsGit/go-ethereum@v1.6.5/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  	stopped  bool
    57  	active   map[ID]*Subscription
    58  	inactive map[ID]*Subscription
    59  }
    60  
    61  // newNotifier creates a new notifier that can be used to send subscription
    62  // notifications to the client.
    63  func newNotifier(codec ServerCodec) *Notifier {
    64  	return &Notifier{
    65  		codec:    codec,
    66  		active:   make(map[ID]*Subscription),
    67  		inactive: make(map[ID]*Subscription),
    68  	}
    69  }
    70  
    71  // NotifierFromContext returns the Notifier value stored in ctx, if any.
    72  func NotifierFromContext(ctx context.Context) (*Notifier, bool) {
    73  	n, ok := ctx.Value(notifierKey{}).(*Notifier)
    74  	return n, ok
    75  }
    76  
    77  // CreateSubscription returns a new subscription that is coupled to the
    78  // RPC connection. By default subscriptions are inactive and notifications
    79  // are dropped until the subscription is marked as active. This is done
    80  // by the RPC server after the subscription ID is send to the client.
    81  func (n *Notifier) CreateSubscription() *Subscription {
    82  	s := &Subscription{ID: NewID(), err: make(chan error)}
    83  	n.subMu.Lock()
    84  	n.inactive[s.ID] = s
    85  	n.subMu.Unlock()
    86  	return s
    87  }
    88  
    89  // Notify sends a notification to the client with the given data as payload.
    90  // If an error occurs the RPC connection is closed and the error is returned.
    91  func (n *Notifier) Notify(id ID, data interface{}) error {
    92  	n.subMu.RLock()
    93  	defer n.subMu.RUnlock()
    94  
    95  	sub, active := n.active[id]
    96  	if active {
    97  		notification := n.codec.CreateNotification(string(id), sub.namespace, data)
    98  		if err := n.codec.Write(notification); err != nil {
    99  			n.codec.Close()
   100  			return err
   101  		}
   102  	}
   103  	return nil
   104  }
   105  
   106  // Closed returns a channel that is closed when the RPC connection is closed.
   107  func (n *Notifier) Closed() <-chan interface{} {
   108  	return n.codec.Closed()
   109  }
   110  
   111  // unsubscribe a subscription.
   112  // If the subscription could not be found ErrSubscriptionNotFound is returned.
   113  func (n *Notifier) unsubscribe(id ID) error {
   114  	n.subMu.Lock()
   115  	defer n.subMu.Unlock()
   116  	if s, found := n.active[id]; found {
   117  		close(s.err)
   118  		delete(n.active, id)
   119  		return nil
   120  	}
   121  	return ErrSubscriptionNotFound
   122  }
   123  
   124  // activate enables a subscription. Until a subscription is enabled all
   125  // notifications are dropped. This method is called by the RPC server after
   126  // the subscription ID was sent to client. This prevents notifications being
   127  // send to the client before the subscription ID is send to the client.
   128  func (n *Notifier) activate(id ID, namespace string) {
   129  	n.subMu.Lock()
   130  	defer n.subMu.Unlock()
   131  	if sub, found := n.inactive[id]; found {
   132  		sub.namespace = namespace
   133  		n.active[id] = sub
   134  		delete(n.inactive, id)
   135  	}
   136  }