github.com/iotexproject/iotex-core@v1.14.1-rc1/action/protocol/registry.go (about)

     1  // Copyright (c) 2019 IoTeX Foundation
     2  // This source code is provided 'as is' and no warranties are given as to title or non-infringement, merchantability
     3  // or fitness for purpose and, to the extent permitted by law, all liability for your use of the code is disclaimed.
     4  // This source code is governed by Apache License 2.0 that can be found in the LICENSE file.
     5  
     6  package protocol
     7  
     8  import (
     9  	"context"
    10  	"reflect"
    11  	"sync"
    12  
    13  	"github.com/pkg/errors"
    14  )
    15  
    16  // Registry is the hub of all protocols deployed on the chain
    17  type Registry struct {
    18  	mu        sync.RWMutex
    19  	ids       map[string]int
    20  	protocols []Protocol
    21  }
    22  
    23  // NewRegistry create a new Registry
    24  func NewRegistry() *Registry {
    25  	return &Registry{
    26  		ids:       make(map[string]int, 0),
    27  		protocols: make([]Protocol, 0),
    28  	}
    29  }
    30  
    31  func (r *Registry) register(id string, p Protocol, force bool) error {
    32  	idx, loaded := r.ids[id]
    33  	if loaded {
    34  		if !force {
    35  			return errors.Errorf("Protocol with ID %s is already registered", id)
    36  		}
    37  		r.protocols[idx] = p
    38  
    39  		return nil
    40  	}
    41  	r.ids[id] = len(r.ids)
    42  	r.protocols = append(r.protocols, p)
    43  
    44  	return nil
    45  }
    46  
    47  // Register registers the protocol with a unique ID
    48  func (r *Registry) Register(id string, p Protocol) error {
    49  	r.mu.Lock()
    50  	defer r.mu.Unlock()
    51  
    52  	return r.register(id, p, false)
    53  }
    54  
    55  // ForceRegister registers the protocol with a unique ID and force replacing the previous protocol if it exists
    56  func (r *Registry) ForceRegister(id string, p Protocol) error {
    57  	r.mu.Lock()
    58  	defer r.mu.Unlock()
    59  
    60  	return r.register(id, p, true)
    61  }
    62  
    63  // Find finds a protocol by ID
    64  func (r *Registry) Find(id string) (Protocol, bool) {
    65  	r.mu.RLock()
    66  	defer r.mu.RUnlock()
    67  	idx, loaded := r.ids[id]
    68  	if !loaded {
    69  		return nil, false
    70  	}
    71  
    72  	return r.protocols[idx], true
    73  }
    74  
    75  // All returns all protocols
    76  func (r *Registry) All() []Protocol {
    77  	if r == nil {
    78  		return nil
    79  	}
    80  	r.mu.RLock()
    81  	defer r.mu.RUnlock()
    82  	return r.all()
    83  }
    84  
    85  func (r *Registry) all() []Protocol {
    86  	all := make([]Protocol, len(r.protocols))
    87  	copy(all, r.protocols)
    88  
    89  	return all
    90  }
    91  
    92  // StartAll starts all protocols which are startable
    93  func (r *Registry) StartAll(ctx context.Context, sr StateReader) (View, error) {
    94  	if r == nil {
    95  		return nil, nil
    96  	}
    97  	r.mu.RLock()
    98  	defer r.mu.RUnlock()
    99  	allView := make(View)
   100  	for _, p := range r.all() {
   101  		s, ok := p.(Starter)
   102  		if !ok {
   103  			continue
   104  		}
   105  		view, err := s.Start(ctx, sr)
   106  		if err != nil {
   107  			return nil, errors.Wrapf(err, "failed to start protocol %s", reflect.TypeOf(p))
   108  		}
   109  		if view == nil {
   110  			continue
   111  		}
   112  		allView[p.Name()] = view
   113  	}
   114  	return allView, nil
   115  }