75 lines
1.7 KiB
Go
75 lines
1.7 KiB
Go
package stream
|
|
|
|
import (
|
|
"context"
|
|
"github.com/anytypeio/go-anytype-infrastructure-experiments/consensus"
|
|
"github.com/cheggaaa/mb/v2"
|
|
"go.uber.org/zap"
|
|
"sync"
|
|
)
|
|
|
|
type Stream struct {
|
|
id uint64
|
|
logIds map[string]struct{}
|
|
mu sync.Mutex
|
|
mb *mb.MB[consensus.Log]
|
|
s *service
|
|
}
|
|
|
|
func (s *Stream) LogIds() [][]byte {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
logIds := make([][]byte, 0, len(s.logIds))
|
|
for logId := range s.logIds {
|
|
logIds = append(logIds, []byte(logId))
|
|
}
|
|
return logIds
|
|
}
|
|
|
|
func (s *Stream) AddRecords(logId []byte, records []consensus.Record) (err error) {
|
|
return s.mb.Add(consensus.Log{Id: logId, Records: records})
|
|
}
|
|
|
|
func (s *Stream) WaitLogs() []consensus.Log {
|
|
return s.mb.Wait()
|
|
}
|
|
|
|
func (s *Stream) WatchIds(ctx context.Context, logIds [][]byte) {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
for _, logId := range logIds {
|
|
logIdKey := string(logId)
|
|
if _, ok := s.logIds[logIdKey]; !ok {
|
|
s.logIds[logIdKey] = struct{}{}
|
|
if addErr := s.s.AddStream(ctx, logId, s); addErr != nil {
|
|
log.Warn("can't add stream for log", zap.Binary("logId", logId), zap.Error(addErr))
|
|
}
|
|
}
|
|
}
|
|
return
|
|
}
|
|
|
|
func (s *Stream) UnwatchIds(ctx context.Context, logIds [][]byte) {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
for _, logId := range logIds {
|
|
logIdKey := string(logId)
|
|
if _, ok := s.logIds[logIdKey]; ok {
|
|
delete(s.logIds, logIdKey)
|
|
if remErr := s.s.RemoveStream(ctx, logId, s.id); remErr != nil {
|
|
log.Warn("can't remove stream for log", zap.Binary("logId", logId), zap.Error(remErr))
|
|
}
|
|
}
|
|
}
|
|
return
|
|
}
|
|
|
|
func (s *Stream) Close() {
|
|
_ = s.mb.Close()
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
for logId := range s.logIds {
|
|
_ = s.s.RemoveStream(context.TODO(), []byte(logId), s.id)
|
|
}
|
|
}
|