github.com/turingchain2020/turingchain@v1.1.21/system/p2p/dht/extension/pubsub.go (about)

     1  package extension
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"sync"
     7  
     8  	"github.com/turingchain2020/turingchain/common/log/log15"
     9  	"github.com/libp2p/go-libp2p-core/host"
    10  	"github.com/libp2p/go-libp2p-core/peer"
    11  	pubsub "github.com/libp2p/go-libp2p-pubsub"
    12  )
    13  
    14  var log = log15.New("module", "pubsub")
    15  
    16  // TopicMap topic map
    17  type TopicMap map[string]*topicinfo
    18  
    19  type topicinfo struct {
    20  	pubtopic *pubsub.Topic
    21  	sub      *pubsub.Subscription
    22  	ctx      context.Context
    23  	cancel   context.CancelFunc
    24  	topic    string
    25  }
    26  
    27  // PubSub pub sub
    28  type PubSub struct {
    29  	ps         *pubsub.PubSub
    30  	topics     TopicMap
    31  	topicMutex sync.RWMutex
    32  	ctx        context.Context
    33  }
    34  
    35  // SubMsg sub message
    36  type SubMsg *pubsub.Message
    37  
    38  // SubCallBack 订阅消息回调函数
    39  type SubCallBack func(topic string, msg SubMsg)
    40  
    41  // NewPubSub new pub sub
    42  func NewPubSub(ctx context.Context, host host.Host, opts ...pubsub.Option) (*PubSub, error) {
    43  	p := &PubSub{
    44  		ps:     nil,
    45  		topics: make(TopicMap),
    46  	}
    47  	//选择使用GossipSub
    48  	ps, err := pubsub.NewGossipSub(ctx, host, opts...)
    49  	if err != nil {
    50  		return nil, err
    51  	}
    52  
    53  	p.ps = ps
    54  	p.ctx = ctx
    55  	p.topics = make(TopicMap)
    56  	return p, nil
    57  }
    58  
    59  // GetTopics get topics
    60  func (p *PubSub) GetTopics() []string {
    61  	return p.ps.GetTopics()
    62  }
    63  
    64  // HasTopic check topic exist
    65  func (p *PubSub) HasTopic(topic string) bool {
    66  	p.topicMutex.RLock()
    67  	defer p.topicMutex.RUnlock()
    68  	_, ok := p.topics[topic]
    69  	return ok
    70  }
    71  
    72  // JoinAndSubTopic 加入topic&subTopic
    73  func (p *PubSub) JoinAndSubTopic(topic string, callback SubCallBack, opts ...pubsub.TopicOpt) error {
    74  
    75  	Topic, err := p.ps.Join(topic, opts...)
    76  	if err != nil {
    77  		return err
    78  	}
    79  
    80  	subscription, err := Topic.Subscribe()
    81  	if err != nil {
    82  		return err
    83  	}
    84  	//p.topics = append(p.topics, Topic)
    85  	ctx, cancel := context.WithCancel(p.ctx)
    86  
    87  	p.topicMutex.Lock()
    88  	p.topics[topic] = &topicinfo{
    89  		pubtopic: Topic,
    90  		ctx:      ctx,
    91  		topic:    topic,
    92  		cancel:   cancel,
    93  		sub:      subscription,
    94  	}
    95  	p.topicMutex.Unlock()
    96  	go p.subTopic(ctx, subscription, callback)
    97  	return nil
    98  }
    99  
   100  // Publish 发布消息
   101  func (p *PubSub) Publish(topic string, msg []byte) error {
   102  	p.topicMutex.RLock()
   103  	defer p.topicMutex.RUnlock()
   104  	t, ok := p.topics[topic]
   105  	if !ok {
   106  		log.Error("pubsub publish", "no this topic", topic)
   107  		return fmt.Errorf("no this topic:%v", topic)
   108  	}
   109  	err := t.pubtopic.Publish(t.ctx, msg)
   110  	if err != nil {
   111  		log.Error("pubsub publish", "err", err)
   112  		return err
   113  	}
   114  	return nil
   115  
   116  }
   117  
   118  func (p *PubSub) subTopic(ctx context.Context, sub *pubsub.Subscription, callback SubCallBack) {
   119  	topic := sub.Topic()
   120  	for {
   121  		got, err := sub.Next(ctx)
   122  		if err != nil {
   123  			log.Error("SubMsg", "topic", topic, "sub err", err)
   124  			p.RemoveTopic(topic)
   125  			return
   126  		}
   127  		callback(topic, got)
   128  	}
   129  }
   130  
   131  // RemoveTopic remove topic
   132  func (p *PubSub) RemoveTopic(topic string) {
   133  
   134  	p.topicMutex.Lock()
   135  	defer p.topicMutex.Unlock()
   136  
   137  	info, ok := p.topics[topic]
   138  	if ok {
   139  		log.Info("RemoveTopic", "topic", topic)
   140  		info.cancel()
   141  		info.sub.Cancel()
   142  		err := info.pubtopic.Close()
   143  		if err != nil {
   144  			log.Error("RemoveTopic", "topic", topic, "close topic err", err)
   145  		}
   146  		delete(p.topics, topic)
   147  	}
   148  
   149  }
   150  
   151  // FetchTopicPeers fetch peers with topic
   152  func (p *PubSub) FetchTopicPeers(topic string) []peer.ID {
   153  	p.topicMutex.RLock()
   154  	defer p.topicMutex.RUnlock()
   155  	topicobj, ok := p.topics[topic]
   156  	if ok {
   157  		return topicobj.pubtopic.ListPeers()
   158  	}
   159  	return nil
   160  }
   161  
   162  // TopicNum get topic number
   163  func (p *PubSub) TopicNum() int {
   164  	return len(p.ps.GetTopics())
   165  
   166  }