github.com/amazechain/amc@v0.1.3/internal/pubsub/pubsub.go (about)

     1  // Copyright 2022 The AmazeChain Authors
     2  // This file is part of the AmazeChain library.
     3  //
     4  // The AmazeChain 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 AmazeChain 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 AmazeChain library. If not, see <http://www.gnu.org/licenses/>.
    16  
    17  package pubsub
    18  
    19  import (
    20  	"context"
    21  	"errors"
    22  	"google.golang.org/protobuf/proto"
    23  	"sync"
    24  	"sync/atomic"
    25  
    26  	"github.com/amazechain/amc/common"
    27  	"github.com/amazechain/amc/common/message"
    28  	"github.com/amazechain/amc/log"
    29  	pubsub "github.com/libp2p/go-libp2p-pubsub"
    30  	"github.com/libp2p/go-libp2p/core/host"
    31  )
    32  
    33  var (
    34  	errorInvalidTopic    = errors.New("invalid topic")
    35  	errorNotRunning      = errors.New("amc pubsub not run")
    36  	errorPubSubIsRunning = errors.New("amc pubsub is running")
    37  )
    38  
    39  type AmcPubSub struct {
    40  	topicLock sync.Mutex
    41  	topicsMap map[string]*pubsub.Topic
    42  
    43  	p2pserver common.INetwork
    44  
    45  	pubsub  *pubsub.PubSub
    46  	running int32
    47  
    48  	host host.Host
    49  
    50  	ctx context.Context
    51  
    52  	chainID uint64
    53  }
    54  
    55  func NewPubSub(ctx context.Context, p2pserver common.INetwork, chainid uint64) (common.IPubSub, error) {
    56  	amc := AmcPubSub{
    57  		ctx:       ctx,
    58  		host:      p2pserver.Host(),
    59  		p2pserver: p2pserver,
    60  		running:   0,
    61  		topicsMap: make(map[string]*pubsub.Topic),
    62  		chainID:   chainid,
    63  	}
    64  
    65  	return &amc, nil
    66  }
    67  
    68  func (m *AmcPubSub) Start() error {
    69  	if m.isRunning() {
    70  		return errorPubSubIsRunning
    71  	}
    72  
    73  	atomic.StoreInt32(&m.running, 1)
    74  
    75  	var options []pubsub.Option
    76  
    77  	options = append(options, pubsub.WithRawTracer(newRawTracer()) /*, pubsub.WithMessageSignaturePolicy(pubsub.MessageSignaturePolicy(0))*/)
    78  	// todo for test
    79  	if false {
    80  		tracer, err := pubsub.NewJSONTracer("./trace.json")
    81  		if err != nil {
    82  			return err
    83  		}
    84  		options = append(options, pubsub.WithEventTracer(tracer))
    85  	}
    86  
    87  	gossip, err := pubsub.NewGossipSub(m.ctx, m.host, options...)
    88  	if err != nil {
    89  		atomic.StoreInt32(&m.running, 0)
    90  		return err
    91  	}
    92  
    93  	gossip.GetTopics()
    94  
    95  	m.pubsub = gossip
    96  
    97  	return nil
    98  }
    99  
   100  func (m *AmcPubSub) JoinTopic(topic string) (*pubsub.Topic, error) {
   101  	if !m.isRunning() {
   102  		return nil, errorNotRunning
   103  	}
   104  	m.topicLock.Lock()
   105  	defer m.topicLock.Unlock()
   106  	if t, ok := m.topicsMap[topic]; ok {
   107  		return t, nil
   108  	}
   109  
   110  	if _, ok := message.TopicMappings[topic]; ok {
   111  		topicHandle, err := m.pubsub.Join(topic)
   112  		if err != nil {
   113  			return nil, err
   114  		}
   115  		m.topicsMap[topic] = topicHandle
   116  		return topicHandle, nil
   117  	}
   118  
   119  	return nil, errorInvalidTopic
   120  }
   121  
   122  func (m *AmcPubSub) isRunning() bool {
   123  	if atomic.LoadInt32(&m.running) <= 0 {
   124  		return false
   125  	}
   126  	return true
   127  }
   128  
   129  func (m *AmcPubSub) Publish(topic string, msg proto.Message) error {
   130  	if !m.isRunning() {
   131  		return errorNotRunning
   132  	}
   133  	m.topicLock.Lock()
   134  	defer m.topicLock.Unlock()
   135  	if t, ok := m.topicsMap[topic]; ok {
   136  		data, err := proto.Marshal(msg)
   137  		if err != nil {
   138  			log.Errorf("failed to publish topic(%s), data: %s, err: %v", topic, msg, err)
   139  			return err
   140  		}
   141  
   142  		return t.Publish(m.ctx, data)
   143  	}
   144  
   145  	return errorInvalidTopic
   146  }
   147  
   148  func (m *AmcPubSub) Subscription(topic string) (*pubsub.Subscription, error) {
   149  	if !m.isRunning() {
   150  		return nil, errorNotRunning
   151  	}
   152  	m.topicLock.Lock()
   153  	defer m.topicLock.Unlock()
   154  	if t, ok := m.topicsMap[topic]; ok {
   155  		return t.Subscribe()
   156  	}
   157  
   158  	return nil, errorInvalidTopic
   159  }
   160  
   161  func (m *AmcPubSub) GetTopics() []string {
   162  	var topics []string
   163  	for k, _ := range m.topicsMap {
   164  		topics = append(topics, k)
   165  	}
   166  
   167  	return topics
   168  }