github.com/Rookout/GoSDK@v0.1.48/pkg/com_ws/size_limited_channel.go (about)

     1  package com_ws
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  	"time"
     7  
     8  	"github.com/Rookout/GoSDK/pkg/config"
     9  	"github.com/Rookout/GoSDK/pkg/rookoutErrors"
    10  )
    11  
    12  
    13  
    14  type SizeLimitedChannel struct {
    15  	channel        chan []byte
    16  	bytesInChannel int
    17  	channelLock    sync.Mutex
    18  	doneChannel    chan []struct{}
    19  	flushing       bool
    20  }
    21  
    22  func NewSizeLimitedChannel() *SizeLimitedChannel {
    23  	return &SizeLimitedChannel{
    24  		channel:     make(chan []byte, config.SizeLimitedChannelConfig().MaxQueueLength),
    25  		doneChannel: make(chan []struct{}, 1),
    26  	}
    27  }
    28  func (s *SizeLimitedChannel) Offer(message []byte) rookoutErrors.RookoutError {
    29  	s.channelLock.Lock()
    30  	defer s.channelLock.Unlock()
    31  
    32  	maxMessageSize := config.SizeLimitedChannelConfig().MaxMessageSize
    33  	if len(message) > maxMessageSize {
    34  		
    35  		
    36  		return rookoutErrors.NewRookMessageSizeExceeded(len(message), maxMessageSize)
    37  	}
    38  	if s.bytesInChannel+len(message) > config.SizeLimitedChannelConfig().MaxBytesInChannel {
    39  		return rookoutErrors.NewRookOutputQueueFull()
    40  	}
    41  
    42  	select {
    43  	case s.channel <- message:
    44  		s.bytesInChannel += len(message)
    45  		return nil
    46  	default:
    47  		return rookoutErrors.NewRookOutputQueueFull()
    48  	}
    49  }
    50  
    51  func (s *SizeLimitedChannel) Poll(ctx context.Context) []byte {
    52  	select {
    53  	case message := <-s.channel:
    54  		s.channelLock.Lock()
    55  		defer s.channelLock.Unlock()
    56  
    57  		s.bytesInChannel -= len(message)
    58  
    59  		if s.bytesInChannel == 0 && s.flushing {
    60  			select {
    61  			case s.doneChannel <- nil:
    62  			default:
    63  			}
    64  		}
    65  
    66  		return message
    67  	case <-ctx.Done():
    68  		return nil
    69  	}
    70  }
    71  
    72  func (s *SizeLimitedChannel) setFlushing(state bool) {
    73  	s.channelLock.Lock()
    74  	defer s.channelLock.Unlock()
    75  	s.flushing = state
    76  }
    77  
    78  func (s *SizeLimitedChannel) Flush() rookoutErrors.RookoutError {
    79  	if s.bytesInChannel == 0 {
    80  		return nil
    81  	}
    82  
    83  	s.setFlushing(true)
    84  	defer func() { s.setFlushing(false) }()
    85  
    86  	timeout := config.SizeLimitedChannelConfig().FlushTimeout
    87  
    88  	select {
    89  	case <-s.doneChannel:
    90  		return nil
    91  	case <-time.After(timeout):
    92  		return rookoutErrors.NewFlushTimedOut()
    93  	}
    94  }