github.com/ydb-platform/ydb-go-sdk/v3@v3.89.2/internal/topic/topicreadercommon/partition_session_storage.go (about)

     1  package topicreadercommon
     2  
     3  import (
     4  	"fmt"
     5  	"sync"
     6  	"time"
     7  
     8  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/grpcwrapper/rawtopic/rawtopicreader"
     9  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
    10  )
    11  
    12  const (
    13  	compactionIntervalTime    = time.Hour
    14  	compactionIntervalRemoves = 10000
    15  )
    16  
    17  type PartitionSessionStorage struct {
    18  	m sync.RWMutex
    19  
    20  	sessions map[rawtopicreader.PartitionSessionID]*sessionInfo
    21  
    22  	removeIndex              int
    23  	lastCompactedTime        time.Time
    24  	lastCompactedRemoveIndex int
    25  }
    26  
    27  func (c *PartitionSessionStorage) initNeedLock() {
    28  	if c.sessions == nil {
    29  		c.sessions = make(map[rawtopicreader.PartitionSessionID]*sessionInfo)
    30  		c.lastCompactedTime = time.Now()
    31  	}
    32  }
    33  
    34  func (c *PartitionSessionStorage) Add(session *PartitionSession) error {
    35  	c.m.Lock()
    36  	defer c.m.Unlock()
    37  
    38  	c.initNeedLock()
    39  
    40  	if _, ok := c.sessions[session.StreamPartitionSessionID]; ok {
    41  		return xerrors.WithStackTrace(fmt.Errorf("session id already existed: %v", session.StreamPartitionSessionID))
    42  	}
    43  	c.sessions[session.StreamPartitionSessionID] = &sessionInfo{Session: session}
    44  
    45  	return nil
    46  }
    47  
    48  func (c *PartitionSessionStorage) Get(id rawtopicreader.PartitionSessionID) (*PartitionSession, error) {
    49  	c.m.RLock()
    50  	defer c.m.RUnlock()
    51  
    52  	c.initNeedLock()
    53  
    54  	partitionInfo, has := c.sessions[id]
    55  	if !has || partitionInfo.Session == nil {
    56  		return nil, xerrors.WithStackTrace(fmt.Errorf("ydb: read undefined partition session with id: %v", id))
    57  	}
    58  
    59  	return partitionInfo.Session, nil
    60  }
    61  
    62  func (c *PartitionSessionStorage) GetAll() []*PartitionSession {
    63  	c.m.Lock()
    64  	defer c.m.Unlock()
    65  
    66  	res := make([]*PartitionSession, 0, len(c.sessions))
    67  	for _, s := range c.sessions {
    68  		res = append(res, s.Session)
    69  	}
    70  
    71  	return res
    72  }
    73  
    74  func (c *PartitionSessionStorage) Remove(id rawtopicreader.PartitionSessionID) (*PartitionSession, error) {
    75  	now := time.Now()
    76  	c.m.Lock()
    77  	defer c.m.Unlock()
    78  
    79  	c.initNeedLock()
    80  
    81  	c.removeIndex++
    82  	if partitionInfo, ok := c.sessions[id]; ok {
    83  		partitionInfo.RemoveTime = now
    84  
    85  		return partitionInfo.Session, nil
    86  	}
    87  
    88  	c.compactionNeedLock(now)
    89  
    90  	return nil, xerrors.WithStackTrace(fmt.Errorf("ydb: delete undefined partition session with id: %v", id))
    91  }
    92  
    93  func (c *PartitionSessionStorage) compactionNeedLock(now time.Time) {
    94  	if !c.isNeedCompactionNeedLock(now) {
    95  		return
    96  	}
    97  	c.doCompactionNeedLock(now)
    98  }
    99  
   100  func (c *PartitionSessionStorage) isNeedCompactionNeedLock(now time.Time) bool {
   101  	return c.removeIndex-c.lastCompactedRemoveIndex < compactionIntervalRemoves &&
   102  		now.Sub(c.lastCompactedTime) < compactionIntervalTime
   103  }
   104  
   105  func (c *PartitionSessionStorage) doCompactionNeedLock(now time.Time) {
   106  	newSessions := make(map[rawtopicreader.PartitionSessionID]*sessionInfo, len(c.sessions))
   107  	for sessionID, info := range c.sessions {
   108  		if info.IsGarbage(c.removeIndex, now) {
   109  			continue
   110  		}
   111  		newSessions[sessionID] = info
   112  	}
   113  	c.sessions = newSessions
   114  }
   115  
   116  type sessionInfo struct {
   117  	RemoveTime   time.Time
   118  	RemovedIndex int
   119  	Session      *PartitionSession
   120  }
   121  
   122  func (si *sessionInfo) IsGarbage(removeIndexNow int, timeNow time.Time) bool {
   123  	return removeIndexNow-si.RemovedIndex >= compactionIntervalRemoves ||
   124  		timeNow.Sub(si.RemoveTime) >= compactionIntervalTime
   125  }