github.com/nsqio/nsq@v1.3.0/apps/nsq_to_file/topic_discoverer.go (about)

     1  package main
     2  
     3  import (
     4  	"os"
     5  	"regexp"
     6  	"sync"
     7  	"time"
     8  
     9  	"github.com/nsqio/go-nsq"
    10  	"github.com/nsqio/nsq/internal/clusterinfo"
    11  	"github.com/nsqio/nsq/internal/http_api"
    12  	"github.com/nsqio/nsq/internal/lg"
    13  )
    14  
    15  type TopicDiscoverer struct {
    16  	logf     lg.AppLogFunc
    17  	opts     *Options
    18  	ci       *clusterinfo.ClusterInfo
    19  	topics   map[string]*FileLogger
    20  	hupChan  chan os.Signal
    21  	termChan chan os.Signal
    22  	wg       sync.WaitGroup
    23  	cfg      *nsq.Config
    24  }
    25  
    26  func newTopicDiscoverer(logf lg.AppLogFunc, opts *Options, cfg *nsq.Config, hupChan chan os.Signal, termChan chan os.Signal) *TopicDiscoverer {
    27  	client := http_api.NewClient(nil, opts.HTTPClientConnectTimeout, opts.HTTPClientRequestTimeout)
    28  	return &TopicDiscoverer{
    29  		logf:     logf,
    30  		opts:     opts,
    31  		ci:       clusterinfo.New(nil, client),
    32  		topics:   make(map[string]*FileLogger),
    33  		hupChan:  hupChan,
    34  		termChan: termChan,
    35  		cfg:      cfg,
    36  	}
    37  }
    38  
    39  func (t *TopicDiscoverer) updateTopics(topics []string) {
    40  	for _, topic := range topics {
    41  		if _, ok := t.topics[topic]; ok {
    42  			continue
    43  		}
    44  
    45  		if !t.isTopicAllowed(topic) {
    46  			t.logf(lg.WARN, "skipping topic %s (doesn't match pattern %s)", topic, t.opts.TopicPattern)
    47  			continue
    48  		}
    49  
    50  		fl, err := NewFileLogger(t.logf, t.opts, topic, t.cfg)
    51  		if err != nil {
    52  			t.logf(lg.ERROR, "couldn't create logger for new topic %s: %s", topic, err)
    53  			continue
    54  		}
    55  		t.topics[topic] = fl
    56  
    57  		t.wg.Add(1)
    58  		go func(fl *FileLogger) {
    59  			fl.router()
    60  			t.wg.Done()
    61  		}(fl)
    62  	}
    63  }
    64  
    65  func (t *TopicDiscoverer) run() {
    66  	var ticker <-chan time.Time
    67  	if len(t.opts.Topics) == 0 {
    68  		ticker = time.Tick(t.opts.TopicRefreshInterval)
    69  	}
    70  	t.updateTopics(t.opts.Topics)
    71  forloop:
    72  	for {
    73  		select {
    74  		case <-ticker:
    75  			newTopics, err := t.ci.GetLookupdTopics(t.opts.NSQLookupdHTTPAddrs)
    76  			if err != nil {
    77  				t.logf(lg.ERROR, "could not retrieve topic list: %s", err)
    78  				continue
    79  			}
    80  			t.updateTopics(newTopics)
    81  		case <-t.termChan:
    82  			for _, fl := range t.topics {
    83  				close(fl.termChan)
    84  			}
    85  			break forloop
    86  		case <-t.hupChan:
    87  			for _, fl := range t.topics {
    88  				fl.hupChan <- true
    89  			}
    90  		}
    91  	}
    92  	t.wg.Wait()
    93  }
    94  
    95  func (t *TopicDiscoverer) isTopicAllowed(topic string) bool {
    96  	if t.opts.TopicPattern == "" {
    97  		return true
    98  	}
    99  	match, err := regexp.MatchString(t.opts.TopicPattern, topic)
   100  	if err != nil {
   101  		return false
   102  	}
   103  	return match
   104  }