github.com/ydb-platform/ydb-go-sdk/v3@v3.57.0/internal/topic/topicreaderinternal/partition_session.go (about)

     1  package topicreaderinternal
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"sync"
     7  	"sync/atomic"
     8  	"time"
     9  
    10  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/grpcwrapper/rawtopic/rawtopicreader"
    11  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xcontext"
    12  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
    13  )
    14  
    15  const (
    16  	compactionIntervalTime    = time.Hour
    17  	compactionIntervalRemoves = 10000
    18  )
    19  
    20  type partitionSession struct {
    21  	Topic       string
    22  	PartitionID int64
    23  
    24  	readerID     int64
    25  	connectionID string
    26  
    27  	ctx                context.Context
    28  	ctxCancel          context.CancelFunc
    29  	partitionSessionID rawtopicreader.PartitionSessionID
    30  
    31  	lastReceivedOffsetEndVal atomic.Int64
    32  	committedOffsetVal       atomic.Int64
    33  }
    34  
    35  func newPartitionSession(
    36  	partitionContext context.Context,
    37  	topic string,
    38  	partitionID int64,
    39  	readerID int64,
    40  	connectionID string,
    41  	partitionSessionID rawtopicreader.PartitionSessionID,
    42  	committedOffset rawtopicreader.Offset,
    43  ) *partitionSession {
    44  	partitionContext, cancel := xcontext.WithCancel(partitionContext)
    45  
    46  	res := &partitionSession{
    47  		Topic:              topic,
    48  		PartitionID:        partitionID,
    49  		readerID:           readerID,
    50  		connectionID:       connectionID,
    51  		ctx:                partitionContext,
    52  		ctxCancel:          cancel,
    53  		partitionSessionID: partitionSessionID,
    54  	}
    55  	res.committedOffsetVal.Store(committedOffset.ToInt64())
    56  	res.lastReceivedOffsetEndVal.Store(committedOffset.ToInt64() - 1)
    57  
    58  	return res
    59  }
    60  
    61  func (s *partitionSession) Context() context.Context {
    62  	return s.ctx
    63  }
    64  
    65  func (s *partitionSession) Close() {
    66  	s.ctxCancel()
    67  }
    68  
    69  func (s *partitionSession) committedOffset() rawtopicreader.Offset {
    70  	v := s.committedOffsetVal.Load()
    71  
    72  	var res rawtopicreader.Offset
    73  	res.FromInt64(v)
    74  
    75  	return res
    76  }
    77  
    78  func (s *partitionSession) setCommittedOffset(v rawtopicreader.Offset) {
    79  	s.committedOffsetVal.Store(v.ToInt64())
    80  }
    81  
    82  func (s *partitionSession) lastReceivedMessageOffset() rawtopicreader.Offset {
    83  	v := s.lastReceivedOffsetEndVal.Load()
    84  
    85  	var res rawtopicreader.Offset
    86  	res.FromInt64(v)
    87  
    88  	return res
    89  }
    90  
    91  func (s *partitionSession) setLastReceivedMessageOffset(v rawtopicreader.Offset) {
    92  	s.lastReceivedOffsetEndVal.Store(v.ToInt64())
    93  }
    94  
    95  type partitionSessionStorage struct {
    96  	m sync.RWMutex
    97  
    98  	sessions map[partitionSessionID]*sessionInfo
    99  
   100  	removeIndex              int
   101  	lastCompactedTime        time.Time
   102  	lastCompactedRemoveIndex int
   103  }
   104  
   105  func (c *partitionSessionStorage) init() {
   106  	c.sessions = make(map[partitionSessionID]*sessionInfo)
   107  	c.lastCompactedTime = time.Now()
   108  }
   109  
   110  func (c *partitionSessionStorage) Add(session *partitionSession) error {
   111  	c.m.Lock()
   112  	defer c.m.Unlock()
   113  
   114  	if _, ok := c.sessions[session.partitionSessionID]; ok {
   115  		return xerrors.WithStackTrace(fmt.Errorf("session id already existed: %v", session.partitionSessionID))
   116  	}
   117  	c.sessions[session.partitionSessionID] = &sessionInfo{Session: session}
   118  
   119  	return nil
   120  }
   121  
   122  func (c *partitionSessionStorage) Get(id partitionSessionID) (*partitionSession, error) {
   123  	c.m.RLock()
   124  	defer c.m.RUnlock()
   125  
   126  	partitionInfo, has := c.sessions[id]
   127  	if !has || partitionInfo.Session == nil {
   128  		return nil, xerrors.WithStackTrace(fmt.Errorf("ydb: read undefined partition session with id: %v", id))
   129  	}
   130  
   131  	return partitionInfo.Session, nil
   132  }
   133  
   134  func (c *partitionSessionStorage) Remove(id partitionSessionID) (*partitionSession, error) {
   135  	now := time.Now()
   136  	c.m.Lock()
   137  	defer c.m.Unlock()
   138  
   139  	c.removeIndex++
   140  	if partitionInfo, ok := c.sessions[id]; ok {
   141  		partitionInfo.RemoveTime = now
   142  
   143  		return partitionInfo.Session, nil
   144  	}
   145  
   146  	c.compactionNeedLock(now)
   147  
   148  	return nil, xerrors.WithStackTrace(fmt.Errorf("ydb: delete undefined partition session with id: %v", id))
   149  }
   150  
   151  func (c *partitionSessionStorage) compactionNeedLock(now time.Time) {
   152  	if !c.isNeedCompactionNeedLock(now) {
   153  		return
   154  	}
   155  	c.doCompactionNeedLock(now)
   156  }
   157  
   158  func (c *partitionSessionStorage) isNeedCompactionNeedLock(now time.Time) bool {
   159  	return c.removeIndex-c.lastCompactedRemoveIndex < compactionIntervalRemoves &&
   160  		now.Sub(c.lastCompactedTime) < compactionIntervalTime
   161  }
   162  
   163  func (c *partitionSessionStorage) doCompactionNeedLock(now time.Time) {
   164  	newSessions := make(map[partitionSessionID]*sessionInfo, len(c.sessions))
   165  	for sessionID, info := range c.sessions {
   166  		if info.IsGarbage(c.removeIndex, now) {
   167  			continue
   168  		}
   169  		newSessions[sessionID] = info
   170  	}
   171  	c.sessions = newSessions
   172  }
   173  
   174  type sessionInfo struct {
   175  	RemoveTime   time.Time
   176  	RemovedIndex int
   177  	Session      *partitionSession
   178  }
   179  
   180  func (si *sessionInfo) IsGarbage(removeIndexNow int, timeNow time.Time) bool {
   181  	return removeIndexNow-si.RemovedIndex >= compactionIntervalRemoves ||
   182  		timeNow.Sub(si.RemoveTime) >= compactionIntervalTime
   183  }