multiclient: assign client an unique id

Author: markya0616

multiclient/client.go

@@ -24,7 +24,7 @@ import (
 	"time"
 
 	"github.com/cskr/pubsub"
-	"github.com/ethereum/go-ethereum"
+	ethereum "github.com/ethereum/go-ethereum"
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/core/types"
 	"github.com/ethereum/go-ethereum/ethclient"
@@ -619,18 +619,18 @@ type Header struct {
 // SubscribeNewHead subscribes to notifications about the current blockchain head
 // on the given channel.
 func (mc *Client) SubscribeNewHead(ctx context.Context, ch chan<- *Header) (ethereum.Subscription, error) {
-	clientsMap := mc.rpcClientMap.Map()
-	lens := len(clientsMap)
+	ids := mc.rpcClientMap.Ids()
+	lens := len(ids)
 	if lens == 0 {
 		return nil, ErrNoEthClient
 	}
 
 	var subscribeNewHeadWg sync.WaitGroup
 
 	cctx, cancel := context.WithCancel(ctx)
-	for url := range clientsMap {
+	for _, id := range ids {
 		subscribeNewHeadWg.Add(1)
-		go mc.subscribeNewHead(cctx, &subscribeNewHeadWg, url, ch)
+		go mc.subscribeNewHead(cctx, &subscribeNewHeadWg, id, ch)
 	}
 
 	newClientCh := mc.pubSub.Sub(newAvailableClientTopic)
@@ -640,9 +640,9 @@ func (mc *Client) SubscribeNewHead(ctx context.Context, ch chan<- *Header) (ethe
 		for {
 			select {
 			case newC := <-newClientCh:
-				url := newC.(string)
+				id := newC.(uint64)
 				subscribeNewHeadWg.Add(1)
-				go mc.subscribeNewHead(cctx, &subscribeNewHeadWg, url, ch)
+				go mc.subscribeNewHead(cctx, &subscribeNewHeadWg, id, ch)
 			case <-cctx.Done():
 				return
 			}
@@ -657,46 +657,69 @@ func (mc *Client) SubscribeNewHead(ctx context.Context, ch chan<- *Header) (ethe
 	}), nil
 }
 
-func (mc *Client) subscribeNewHead(ctx context.Context, wg *sync.WaitGroup, url string, ch chan<- *Header) error {
+func (mc *Client) subscribeNewHead(ctx context.Context, wg *sync.WaitGroup, id uint64, ch chan<- *Header) error {
+	logger := log.New("id", id)
 	defer wg.Done()
 
 	for {
-		rc := mc.rpcClientMap.Get(url)
-		if rc == nil {
-			log.Trace("EthClient has been removed", "url", url)
+		rc, existed := mc.rpcClientMap.GetById(id)
+		if !existed {
+			logger.Trace("EthClient has been removed")
 			return nil
 		}
-
-		headerCh := make(chan *types.Header)
-		c := ethclient.NewClient(rc)
-		sub, err := c.SubscribeNewHead(ctx, headerCh)
-		if err != nil {
-			log.Warn("Failed to subscribe new head", "url", url, "err", err)
+		if rc == nil {
+			logger.Warn("The client is not ready, wait for a while")
 		} else {
-		WAIT_NEW_HEADER:
-			for {
-				select {
-				case header := <-headerCh:
-					h := &Header{
-						Client: rc,
-						Header: header,
-					}
-					select {
-					case ch <- h:
-					case <-ctx.Done():
-						return nil
-					}
-				case err := <-sub.Err():
-					log.Warn("Failed during subscription", "url", url, "err", err)
-					break WAIT_NEW_HEADER
-				case <-ctx.Done():
-					return nil
-				}
+			// If we have error, we need to retry
+			err := doSubscribe(ctx, id, rc, ch)
+			if err == nil {
+				return nil
 			}
 		}
+
 		// retry subscribe after retryPeriod
-		time.Sleep(retryPeriod)
-		log.Trace("Retry to subscribe new head", "url", url)
+		t := time.NewTicker(retryPeriod)
+		select {
+		case <-t.C:
+		case <-ctx.Done():
+			t.Stop()
+			return nil
+		}
+		t.Stop()
+		logger.Trace("Retry to subscribe new head")
+	}
+}
+
+func doSubscribe(ctx context.Context, id uint64, rc *rpc.Client, ch chan<- *Header) error {
+	logger := log.New("id", id)
+	headerCh := make(chan *types.Header)
+	c := ethclient.NewClient(rc)
+	subCtx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
+	sub, err := c.SubscribeNewHead(subCtx, headerCh)
+	if err != nil {
+		logger.Warn("Failed to subscribe new head", "err", err)
+		return err
+	}
+	for {
+		select {
+		case header := <-headerCh:
+			h := &Header{
+				Client: rc,
+				Header: header,
+			}
+			select {
+			case ch <- h:
+			case <-subCtx.Done():
+				return nil
+			}
+		case err := <-sub.Err():
+			logger.Warn("Failed during subscription", "err", err)
+			return err
+		case <-subCtx.Done():
+			return nil
+		}
 	}
 }
 
@@ -788,8 +811,8 @@ func (mc *Client) DialClients(ctx context.Context) {
 	for i := 0; i < len(urls); i++ {
 		dialed := <-dialCh
 		if dialed.client != nil {
-			mc.rpcClientMap.Replace(dialed.url, dialed.client)
-			mc.pubSub.Pub(dialed.url, newAvailableClientTopic)
+			id := mc.rpcClientMap.Replace(dialed.url, dialed.client)
+			mc.pubSub.Pub(id, newAvailableClientTopic)
 		}
 	}
 }

multiclient/consul.go

@@ -38,7 +38,7 @@ func ConsulDiscovery(rawURL, serviceID, serviceScheme string) Option {
 		}
 		log.Info("EthClients from consul", "urls", urls)
 		for _, url := range urls {
-			mc.ClientMap().Set(url, nil)
+			mc.ClientMap().Add(url, nil)
 		}
 		return nil
 	}

multiclient/k8s.go

@@ -21,7 +21,7 @@ import (
 	"sync"
 
 	"github.com/getamis/sirius/log"
-	"k8s.io/api/core/v1"
+	v1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/meta"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/fields"
@@ -180,7 +180,7 @@ func (s *endpointStore) Add(obj interface{}) error {
 	defer s.mutex.Unlock()
 
 	for _, url := range urls {
-		s.rpcClientMap.Set(url, nil)
+		s.rpcClientMap.Add(url, nil)
 	}
 
 	s.endpoints[o.GetUID()] = urls
@@ -202,7 +202,7 @@ func (s *endpointStore) Update(obj interface{}) error {
 	// Add new urls
 	for _, url := range news {
 		if c := s.rpcClientMap.Get(url); c == nil {
-			s.rpcClientMap.Set(url, nil)
+			s.rpcClientMap.Add(url, nil)
 		}
 	}
 

multiclient/map.go

@@ -24,65 +24,104 @@ import (
 )
 
 type Map struct {
-	m           map[string]*rpc.Client
-	newClientCh chan<- string
+	// the mapping form url to client
+	clientMap map[string]*client
+	// the mapping subscription id to url
+	idMap              map[uint64]string
+	subscripionCounter uint64
+	newClientCh        chan<- string
 
 	lock sync.RWMutex
 }
 
+type client struct {
+	*rpc.Client
+	Id uint64
+}
+
 func NewMap(newClientCh chan<- string) *Map {
 	return &Map{
-		m:           make(map[string]*rpc.Client),
-		newClientCh: newClientCh,
+		clientMap:          make(map[string]*client),
+		idMap:              make(map[uint64]string),
+		subscripionCounter: 0,
+		newClientCh:        newClientCh,
 	}
 }
 
 func (m *Map) Delete(key string) {
 	m.lock.Lock()
 	defer m.lock.Unlock()
 
-	if c := m.m[key]; c != nil {
-		c.Close()
+	c, ok := m.clientMap[key]
+	if !ok {
+		return
+	}
+	if c.Client != nil {
+		c.Client.Close()
 	}
-	delete(m.m, key)
-	log.Trace("Eth client removed", "url", key)
+	delete(m.idMap, c.Id)
+	delete(m.clientMap, key)
+	log.Trace("Eth client removed", "c.Id", c.Id, "url", key)
 }
 
-func (m *Map) Set(key string, value *rpc.Client) {
+func (m *Map) Add(key string, value *rpc.Client) {
 	m.lock.Lock()
 	defer m.lock.Unlock()
 
-	m.m[key] = value
+	m.subscripionCounter++
+	m.clientMap[key] = &client{
+		Id:     m.subscripionCounter,
+		Client: value,
+	}
+	m.idMap[m.subscripionCounter] = key
+
 	if m.newClientCh != nil {
 		select {
 		case m.newClientCh <- key:
 		default:
 		}
 	}
-	log.Trace("Eth client added", "url", key)
+	log.Trace("Eth client added", "id", m.subscripionCounter, "url", key)
 }
 
-func (m *Map) Replace(key string, value *rpc.Client) {
+func (m *Map) Replace(key string, value *rpc.Client) uint64 {
 	m.lock.Lock()
 	defer m.lock.Unlock()
 
-	if _, ok := m.m[key]; ok {
-		m.m[key] = value
+	if _, ok := m.clientMap[key]; ok {
+		m.clientMap[key].Client = value
 	}
+
+	return m.clientMap[key].Id
 }
 
 func (m *Map) Get(key string) *rpc.Client {
 	m.lock.RLock()
 	defer m.lock.RUnlock()
 
-	return m.m[key]
+	return m.clientMap[key].Client
+}
+
+func (m *Map) GetById(id uint64) (*rpc.Client, bool) {
+	m.lock.RLock()
+	defer m.lock.RUnlock()
+
+	key, ok := m.idMap[id]
+	if !ok {
+		return nil, false
+	}
+	c, ok := m.clientMap[key]
+	if !ok {
+		return nil, false
+	}
+	return c.Client, true
 }
 
 func (m *Map) Len() int {
 	m.lock.RLock()
 	defer m.lock.RUnlock()
 
-	return len(m.m)
+	return len(m.clientMap)
 }
 
 // List returns a deep copy of client list
@@ -91,9 +130,9 @@ func (m *Map) List() []*rpc.Client {
 	defer m.lock.RUnlock()
 
 	l := []*rpc.Client{}
-	for _, v := range m.m {
-		if v != nil {
-			l = append(l, v)
+	for _, v := range m.clientMap {
+		if v.Client != nil {
+			l = append(l, v.Client)
 		}
 	}
 	return l
@@ -105,21 +144,34 @@ func (m *Map) Map() map[string]*rpc.Client {
 	defer m.lock.RUnlock()
 
 	newMap := map[string]*rpc.Client{}
-	for k, v := range m.m {
-		if v != nil {
-			newMap[k] = v
+	for k, v := range m.clientMap {
+		if v.Client != nil {
+			newMap[k] = v.Client
 		}
 	}
 	return newMap
 }
 
+func (m *Map) Ids() []uint64 {
+	m.lock.RLock()
+	defer m.lock.RUnlock()
+
+	ids := make([]uint64, len(m.idMap))
+	i := 0
+	for k := range m.idMap {
+		ids[i] = k
+		i++
+	}
+	return ids
+}
+
 func (m *Map) Keys() []string {
 	m.lock.RLock()
 	defer m.lock.RUnlock()
 
-	urls := make([]string, len(m.m))
+	urls := make([]string, len(m.clientMap))
 	index := 0
-	for k := range m.m {
+	for k := range m.clientMap {
 		urls[index] = k
 		index++
 	}
@@ -131,8 +183,8 @@ func (m *Map) NilClients() []string {
 	defer m.lock.RUnlock()
 
 	urls := make([]string, 0)
-	for k, v := range m.m {
-		if v == nil {
+	for k, v := range m.clientMap {
+		if v.Client == nil {
 			urls = append(urls, k)
 		}
 	}

multiclient/option.go

@@ -28,7 +28,7 @@ func EthURLs(urls []string) Option {
 	return func(mc *Client) error {
 		log.Info("EthClients from static list", "urls", urls)
 		for _, url := range urls {
-			mc.ClientMap().Set(url, nil)
+			mc.ClientMap().Add(url, nil)
 		}
 		return nil
 	}