code.vegaprotocol.io/vega@v0.79.0/libs/subscribers/service.go (about)

     1  // Copyright (C) 2023 Gobalsky Labs Limited
     2  //
     3  // This program is free software: you can redistribute it and/or modify
     4  // it under the terms of the GNU Affero General Public License as
     5  // published by the Free Software Foundation, either version 3 of the
     6  // License, or (at your option) any later version.
     7  //
     8  // This program is distributed in the hope that it will be useful,
     9  // but WITHOUT ANY WARRANTY; without even the implied warranty of
    10  // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    11  // GNU Affero General Public License for more details.
    12  //
    13  // You should have received a copy of the GNU Affero General Public License
    14  // along with this program.  If not, see <http://www.gnu.org/licenses/>.
    15  
    16  package subscribers
    17  
    18  import (
    19  	"context"
    20  	"time"
    21  
    22  	"code.vegaprotocol.io/vega/core/events"
    23  	"code.vegaprotocol.io/vega/libs/broker"
    24  	"code.vegaprotocol.io/vega/logging"
    25  	eventspb "code.vegaprotocol.io/vega/protos/vega/events/v1"
    26  )
    27  
    28  //go:generate go run github.com/golang/mock/mockgen -destination mocks/broker_mocks.go -package mocks code.vegaprotocol.io/vega/libs/subscribers Broker
    29  type Broker interface {
    30  	Subscribe(s broker.Subscriber) int
    31  	Unsubscribe(id int)
    32  }
    33  
    34  type Service struct {
    35  	log           *logging.Logger
    36  	broker        Broker
    37  	maxBufferSize int
    38  }
    39  
    40  type StreamSubscription interface {
    41  	Halt()
    42  	Push(evts ...events.Event)
    43  	UpdateBatchSize(ctx context.Context, size int) []*eventspb.BusEvent
    44  	Types() []events.Type
    45  	GetData(ctx context.Context) []*eventspb.BusEvent
    46  	C() chan<- []events.Event
    47  	Closed() <-chan struct{}
    48  	Skip() <-chan struct{}
    49  	SetID(id int)
    50  	ID() int
    51  	Ack() bool
    52  }
    53  
    54  const namedLogger = "subscribers"
    55  
    56  func NewService(log *logging.Logger, broker Broker, maxBufferSize int) *Service {
    57  	return &Service{
    58  		log:           log.Named(namedLogger),
    59  		broker:        broker,
    60  		maxBufferSize: maxBufferSize,
    61  	}
    62  }
    63  
    64  func (s *Service) ObserveEvents(ctx context.Context, retries int, eTypes []events.Type, batchSize int, filters ...EventFilter) (<-chan []*eventspb.BusEvent, chan<- int) {
    65  	return s.ObserveEventsOnStream(ctx, retries, NewStreamSub(ctx, eTypes, batchSize, filters...))
    66  }
    67  
    68  func (s *Service) ObserveEventsOnStream(ctx context.Context, retries int,
    69  	sub StreamSubscription,
    70  ) (<-chan []*eventspb.BusEvent, chan<- int) {
    71  	// one batch buffer for the out channel
    72  	in, out := make(chan int), make(chan []*eventspb.BusEvent, 1)
    73  	ctx, cfunc := context.WithCancel(ctx)
    74  
    75  	// use stream subscriber
    76  	// use buffer size of 0 for the time being
    77  	id := s.broker.Subscribe(sub)
    78  
    79  	// makes the tick duration 2000ms max to wait basically
    80  	tickDuration := 10 * time.Millisecond
    81  	retries = 200
    82  
    83  	go func() {
    84  		data := []*eventspb.BusEvent{}
    85  		defer func() {
    86  			s.broker.Unsubscribe(id)
    87  			close(out)
    88  			cfunc()
    89  		}()
    90  		ret := retries
    91  
    92  		trySend := func() {
    93  			t := time.NewTicker(tickDuration)
    94  			defer t.Stop()
    95  			select {
    96  			case <-ctx.Done():
    97  				return
    98  			case out <- data:
    99  				data = []*eventspb.BusEvent{}
   100  				ret = retries
   101  			case <-t.C:
   102  				if ret == 0 {
   103  					return
   104  				}
   105  				ret--
   106  			}
   107  		}
   108  
   109  		for {
   110  			select {
   111  			case <-ctx.Done():
   112  				return
   113  			case bs := <-in:
   114  				// batch size changed: drain buffer and send data
   115  				data = append(data, sub.UpdateBatchSize(ctx, bs)...)
   116  				dataLength := len(data)
   117  
   118  				if dataLength > s.maxBufferSize {
   119  					s.log.Warningf("slow consumer detected, closing event observer")
   120  					return
   121  				}
   122  
   123  				if dataLength > 0 {
   124  					trySend()
   125  				}
   126  			default:
   127  				// wait for actual changes
   128  				data = append(data, sub.GetData(ctx)...)
   129  				dataLength := len(data)
   130  
   131  				if dataLength > s.maxBufferSize {
   132  					s.log.Warningf("slow consumer detected, closing event observer")
   133  					return
   134  				}
   135  
   136  				// this is a very rare thing, but it can happen
   137  				if dataLength > 0 {
   138  					trySend()
   139  				}
   140  			}
   141  		}
   142  	}()
   143  	return out, in
   144  }