github.com/wfusion/gofusion@v1.1.14/common/infra/asynq/subscriber.go (about)

     1  // Copyright 2020 Kentaro Hibino. All rights reserved.
     2  // Use of this source code is governed by a MIT license
     3  // that can be found in the LICENSE file.
     4  
     5  package asynq
     6  
     7  import (
     8  	"sync"
     9  	"time"
    10  
    11  	"github.com/redis/go-redis/v9"
    12  	"github.com/wfusion/gofusion/common/infra/asynq/pkg/base"
    13  	"github.com/wfusion/gofusion/common/infra/asynq/pkg/log"
    14  )
    15  
    16  type subscriber struct {
    17  	logger *log.Logger
    18  	broker base.Broker
    19  
    20  	// channel to communicate back to the long running "subscriber" goroutine.
    21  	done chan struct{}
    22  
    23  	// cancelations hold cancel functions for all active tasks.
    24  	cancelations *base.Cancelations
    25  
    26  	// time to wait before retrying to connect to redis.
    27  	retryTimeout time.Duration
    28  }
    29  
    30  type subscriberParams struct {
    31  	logger       *log.Logger
    32  	broker       base.Broker
    33  	cancelations *base.Cancelations
    34  }
    35  
    36  func newSubscriber(params subscriberParams) *subscriber {
    37  	return &subscriber{
    38  		logger:       params.logger,
    39  		broker:       params.broker,
    40  		done:         make(chan struct{}),
    41  		cancelations: params.cancelations,
    42  		retryTimeout: 5 * time.Second,
    43  	}
    44  }
    45  
    46  func (s *subscriber) shutdown() {
    47  	s.logger.Debug("[Common] asynq subscriber shutting down...")
    48  	// Signal the subscriber goroutine to stop.
    49  	s.done <- struct{}{}
    50  }
    51  
    52  func (s *subscriber) start(wg *sync.WaitGroup) {
    53  	wg.Add(1)
    54  	go func() {
    55  		defer wg.Done()
    56  		var (
    57  			pubsub *redis.PubSub
    58  			err    error
    59  		)
    60  		// Try until successfully connect to Redis.
    61  		for {
    62  			pubsub, err = s.broker.CancelationPubSub()
    63  			if err != nil {
    64  				s.logger.Errorf("[Common] asynq cannot subscribe to cancelation channel: %v", err)
    65  				select {
    66  				case <-time.After(s.retryTimeout):
    67  					continue
    68  				case <-s.done:
    69  					s.logger.Debug("[Common] asynq subscriber done")
    70  					return
    71  				}
    72  			}
    73  			break
    74  		}
    75  		cancelCh := pubsub.Channel()
    76  		for {
    77  			select {
    78  			case <-s.done:
    79  				_ = pubsub.Close()
    80  				s.logger.Debug("[Common] asynq subscriber done")
    81  				return
    82  			case msg, ok := <-cancelCh:
    83  				if !ok {
    84  					return
    85  				}
    86  				if msg != nil {
    87  					cancel, ok := s.cancelations.Get(msg.Payload)
    88  					if ok {
    89  						cancel()
    90  					}
    91  				}
    92  			}
    93  		}
    94  	}()
    95  }