github.com/prebid/prebid-server/v2@v2.18.0/analytics/pubstack/eventchannel/eventchannel.go (about)

     1  package eventchannel
     2  
     3  import (
     4  	"bytes"
     5  	"compress/gzip"
     6  	"sync"
     7  	"time"
     8  
     9  	"github.com/benbjohnson/clock"
    10  	"github.com/golang/glog"
    11  )
    12  
    13  type Metrics struct {
    14  	bufferSize int64
    15  	eventCount int64
    16  }
    17  
    18  type Limit struct {
    19  	maxByteSize   int64
    20  	maxEventCount int64
    21  	maxTime       time.Duration
    22  }
    23  
    24  type EventChannel struct {
    25  	gz   *gzip.Writer
    26  	buff *bytes.Buffer
    27  
    28  	ch          chan []byte
    29  	endCh       chan int
    30  	metrics     Metrics
    31  	muxGzBuffer sync.RWMutex
    32  	send        Sender
    33  	limit       Limit
    34  	clock       clock.Clock
    35  }
    36  
    37  func NewEventChannel(sender Sender, clock clock.Clock, maxByteSize, maxEventCount int64, maxTime time.Duration) *EventChannel {
    38  	b := &bytes.Buffer{}
    39  	gzw := gzip.NewWriter(b)
    40  
    41  	c := EventChannel{
    42  		gz:      gzw,
    43  		buff:    b,
    44  		ch:      make(chan []byte),
    45  		endCh:   make(chan int),
    46  		metrics: Metrics{},
    47  		send:    sender,
    48  		limit:   Limit{maxByteSize, maxEventCount, maxTime},
    49  		clock:   clock,
    50  	}
    51  	go c.start()
    52  	return &c
    53  }
    54  
    55  func (c *EventChannel) Push(event []byte) {
    56  	c.ch <- event
    57  }
    58  
    59  func (c *EventChannel) Close() {
    60  	c.endCh <- 1
    61  }
    62  
    63  func (c *EventChannel) buffer(event []byte) {
    64  	c.muxGzBuffer.Lock()
    65  	defer c.muxGzBuffer.Unlock()
    66  
    67  	_, err := c.gz.Write(event)
    68  	if err != nil {
    69  		glog.Warning("[pubstack] fail to compress, skip the event")
    70  		return
    71  	}
    72  
    73  	c.metrics.eventCount++
    74  	c.metrics.bufferSize += int64(len(event))
    75  }
    76  
    77  func (c *EventChannel) isBufferFull() bool {
    78  	c.muxGzBuffer.RLock()
    79  	defer c.muxGzBuffer.RUnlock()
    80  	return c.metrics.eventCount >= c.limit.maxEventCount || c.metrics.bufferSize >= c.limit.maxByteSize
    81  }
    82  
    83  func (c *EventChannel) reset() {
    84  	// reset buffer
    85  	c.gz.Reset(c.buff)
    86  	c.buff.Reset()
    87  
    88  	// reset metrics
    89  	c.metrics.eventCount = 0
    90  	c.metrics.bufferSize = 0
    91  }
    92  
    93  func (c *EventChannel) flush() {
    94  	c.muxGzBuffer.Lock()
    95  	defer c.muxGzBuffer.Unlock()
    96  
    97  	if c.metrics.eventCount == 0 || c.metrics.bufferSize == 0 {
    98  		return
    99  	}
   100  
   101  	// reset buffers and writers
   102  	defer c.reset()
   103  
   104  	// finish writing gzip header
   105  	err := c.gz.Close()
   106  	if err != nil {
   107  		glog.Warning("[pubstack] fail to close gzipped buffer")
   108  		return
   109  	}
   110  
   111  	// copy the current buffer to send the payload in a new thread
   112  	payload := make([]byte, c.buff.Len())
   113  	_, err = c.buff.Read(payload)
   114  	if err != nil {
   115  		glog.Warning("[pubstack] fail to copy the buffer")
   116  		return
   117  	}
   118  
   119  	// send events (async)
   120  	go c.send(payload)
   121  }
   122  
   123  func (c *EventChannel) start() {
   124  	ticker := c.clock.Ticker(c.limit.maxTime)
   125  
   126  	for {
   127  		select {
   128  		case <-c.endCh:
   129  			c.flush()
   130  			return
   131  
   132  		// event is received
   133  		case event := <-c.ch:
   134  			c.buffer(event)
   135  			if c.isBufferFull() {
   136  				c.flush()
   137  			}
   138  
   139  		// time between 2 flushes has passed
   140  		case <-ticker.C:
   141  			c.flush()
   142  		}
   143  	}
   144  }