github.com/Jeffail/benthos/v3@v3.65.0/lib/input/reader/redis_pubsub.go (about)

     1  package reader
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  	"time"
     7  
     8  	bredis "github.com/Jeffail/benthos/v3/internal/impl/redis"
     9  	"github.com/Jeffail/benthos/v3/lib/log"
    10  	"github.com/Jeffail/benthos/v3/lib/message"
    11  	"github.com/Jeffail/benthos/v3/lib/metrics"
    12  	"github.com/Jeffail/benthos/v3/lib/types"
    13  	"github.com/go-redis/redis/v7"
    14  )
    15  
    16  //------------------------------------------------------------------------------
    17  
    18  // RedisPubSubConfig contains configuration fields for the RedisPubSub input
    19  // type.
    20  type RedisPubSubConfig struct {
    21  	bredis.Config `json:",inline" yaml:",inline"`
    22  	Channels      []string `json:"channels" yaml:"channels"`
    23  	UsePatterns   bool     `json:"use_patterns" yaml:"use_patterns"`
    24  }
    25  
    26  // NewRedisPubSubConfig creates a new RedisPubSubConfig with default values.
    27  func NewRedisPubSubConfig() RedisPubSubConfig {
    28  	return RedisPubSubConfig{
    29  		Config:      bredis.NewConfig(),
    30  		Channels:    []string{"benthos_chan"},
    31  		UsePatterns: false,
    32  	}
    33  }
    34  
    35  //------------------------------------------------------------------------------
    36  
    37  // RedisPubSub is an input type that reads Redis Pub/Sub messages.
    38  type RedisPubSub struct {
    39  	client redis.UniversalClient
    40  	pubsub *redis.PubSub
    41  	cMut   sync.Mutex
    42  
    43  	conf RedisPubSubConfig
    44  
    45  	stats metrics.Type
    46  	log   log.Modular
    47  }
    48  
    49  // NewRedisPubSub creates a new RedisPubSub input type.
    50  func NewRedisPubSub(
    51  	conf RedisPubSubConfig, log log.Modular, stats metrics.Type,
    52  ) (*RedisPubSub, error) {
    53  	r := &RedisPubSub{
    54  		conf:  conf,
    55  		stats: stats,
    56  		log:   log,
    57  	}
    58  
    59  	_, err := r.conf.Config.Client()
    60  	if err != nil {
    61  		return nil, err
    62  	}
    63  
    64  	return r, nil
    65  }
    66  
    67  //------------------------------------------------------------------------------
    68  
    69  // Connect establishes a connection to a RedisPubSub server.
    70  func (r *RedisPubSub) Connect() error {
    71  	return r.ConnectWithContext(context.Background())
    72  }
    73  
    74  // ConnectWithContext establishes a connection to an RedisPubSub server.
    75  func (r *RedisPubSub) ConnectWithContext(ctx context.Context) error {
    76  	r.cMut.Lock()
    77  	defer r.cMut.Unlock()
    78  
    79  	if r.client != nil {
    80  		return nil
    81  	}
    82  
    83  	client, err := r.conf.Config.Client()
    84  	if err != nil {
    85  		return err
    86  	}
    87  	if _, err := client.Ping().Result(); err != nil {
    88  		return err
    89  	}
    90  
    91  	r.log.Infof("Receiving Redis pub/sub messages from channels: %v\n", r.conf.Channels)
    92  
    93  	r.client = client
    94  	if r.conf.UsePatterns {
    95  		r.pubsub = r.client.PSubscribe(r.conf.Channels...)
    96  	} else {
    97  		r.pubsub = r.client.Subscribe(r.conf.Channels...)
    98  	}
    99  	return nil
   100  }
   101  
   102  // Read attempts to pop a message from a redis pubsub channel.
   103  func (r *RedisPubSub) Read() (types.Message, error) {
   104  	msg, _, err := r.ReadWithContext(context.Background())
   105  	return msg, err
   106  }
   107  
   108  // ReadWithContext attempts to pop a message from a redis pubsub channel.
   109  func (r *RedisPubSub) ReadWithContext(ctx context.Context) (types.Message, AsyncAckFn, error) {
   110  	var pubsub *redis.PubSub
   111  
   112  	r.cMut.Lock()
   113  	pubsub = r.pubsub
   114  	r.cMut.Unlock()
   115  
   116  	if pubsub == nil {
   117  		return nil, nil, types.ErrNotConnected
   118  	}
   119  
   120  	select {
   121  	case rMsg, open := <-pubsub.Channel():
   122  		if !open {
   123  			r.disconnect()
   124  			return nil, nil, types.ErrTypeClosed
   125  		}
   126  		return message.New([][]byte{[]byte(rMsg.Payload)}), noopAsyncAckFn, nil
   127  	case <-ctx.Done():
   128  	}
   129  
   130  	return nil, nil, types.ErrTimeout
   131  }
   132  
   133  // Acknowledge is a noop since Redis pub/sub channels do not support
   134  // acknowledgements.
   135  func (r *RedisPubSub) Acknowledge(err error) error {
   136  	return nil
   137  }
   138  
   139  // disconnect safely closes a connection to an RedisPubSub server.
   140  func (r *RedisPubSub) disconnect() error {
   141  	r.cMut.Lock()
   142  	defer r.cMut.Unlock()
   143  
   144  	var err error
   145  	if r.pubsub != nil {
   146  		err = r.pubsub.Close()
   147  		r.pubsub = nil
   148  	}
   149  	if r.client != nil {
   150  		err = r.client.Close()
   151  		r.client = nil
   152  	}
   153  	return err
   154  }
   155  
   156  // CloseAsync shuts down the RedisPubSub input and stops processing requests.
   157  func (r *RedisPubSub) CloseAsync() {
   158  	r.disconnect()
   159  }
   160  
   161  // WaitForClose blocks until the RedisPubSub input has closed down.
   162  func (r *RedisPubSub) WaitForClose(timeout time.Duration) error {
   163  	return nil
   164  }
   165  
   166  //------------------------------------------------------------------------------