use multiqueue for stream + MessageQueueId interface

This commit is contained in:
Sergey Cherepanov 2023-05-15 19:52:19 +02:00 committed by Mikhail Iudin
parent 34c6ce261e
commit 3d4ac4d130
No known key found for this signature in database
GPG Key ID: FAAAA8BAABDFF1C0
2 changed files with 18 additions and 22 deletions

View File

@ -3,7 +3,7 @@ package streampool
import ( import (
"context" "context"
"github.com/anytypeio/any-sync/app/logger" "github.com/anytypeio/any-sync/app/logger"
"github.com/cheggaaa/mb/v3" "github.com/anytypeio/any-sync/util/multiqueue"
"go.uber.org/zap" "go.uber.org/zap"
"storj.io/drpc" "storj.io/drpc"
"sync/atomic" "sync/atomic"
@ -17,12 +17,16 @@ type stream struct {
streamId uint32 streamId uint32
closed atomic.Bool closed atomic.Bool
l logger.CtxLogger l logger.CtxLogger
queue *mb.MB[drpc.Message] queue multiqueue.MultiQueue[drpc.Message]
tags []string tags []string
} }
func (sr *stream) write(msg drpc.Message) (err error) { func (sr *stream) write(msg drpc.Message) (err error) {
return sr.queue.TryAdd(msg) var queueId string
if qId, ok := msg.(MessageQueueId); ok {
queueId = qId.MessageQueueId()
}
return sr.queue.Add(sr.stream.Context(), queueId, msg)
} }
func (sr *stream) readLoop() error { func (sr *stream) readLoop() error {
@ -45,20 +49,13 @@ func (sr *stream) readLoop() error {
} }
} }
func (sr *stream) writeLoop() error { func (sr *stream) writeToStream(msg drpc.Message) {
defer func() { if err := sr.stream.MsgSend(msg, EncodingProto); err != nil {
sr.l.Warn("msg send error", zap.Error(err))
sr.streamClose() sr.streamClose()
}() return
sr.l.Debug("stream write started")
for {
msg, err := sr.queue.WaitOne(sr.stream.Context())
if err != nil {
return err
}
if err = sr.stream.MsgSend(msg, EncodingProto); err != nil {
return err
}
} }
return
} }
func (sr *stream) streamClose() { func (sr *stream) streamClose() {

View File

@ -4,7 +4,7 @@ import (
"fmt" "fmt"
"github.com/anytypeio/any-sync/net/peer" "github.com/anytypeio/any-sync/net/peer"
"github.com/anytypeio/any-sync/net/pool" "github.com/anytypeio/any-sync/net/pool"
"github.com/cheggaaa/mb/v3" "github.com/anytypeio/any-sync/util/multiqueue"
"go.uber.org/zap" "go.uber.org/zap"
"golang.org/x/exp/slices" "golang.org/x/exp/slices"
"golang.org/x/net/context" "golang.org/x/net/context"
@ -25,6 +25,10 @@ type StreamHandler interface {
// PeerGetter should dial or return cached peers // PeerGetter should dial or return cached peers
type PeerGetter func(ctx context.Context) (peers []peer.Peer, err error) type PeerGetter func(ctx context.Context) (peers []peer.Peer, err error)
type MessageQueueId interface {
MessageQueueId() string
}
// StreamPool keeps and read streams // StreamPool keeps and read streams
type StreamPool interface { type StreamPool interface {
// AddStream adds new outgoing stream into the pool // AddStream adds new outgoing stream into the pool
@ -115,19 +119,14 @@ func (s *streamPool) addStream(drpcStream drpc.Stream, tags ...string) (*stream,
pool: s, pool: s,
streamId: streamId, streamId: streamId,
l: log.With(zap.String("peerId", peerId), zap.Uint32("streamId", streamId)), l: log.With(zap.String("peerId", peerId), zap.Uint32("streamId", streamId)),
queue: mb.New[drpc.Message](queueSize),
tags: tags, tags: tags,
} }
st.queue = multiqueue.New[drpc.Message](st.writeToStream, s.writeQueueSize)
s.streams[streamId] = st s.streams[streamId] = st
s.streamIdsByPeer[peerId] = append(s.streamIdsByPeer[peerId], streamId) s.streamIdsByPeer[peerId] = append(s.streamIdsByPeer[peerId], streamId)
for _, tag := range tags { for _, tag := range tags {
s.streamIdsByTag[tag] = append(s.streamIdsByTag[tag], streamId) s.streamIdsByTag[tag] = append(s.streamIdsByTag[tag], streamId)
} }
go func() {
if err := st.writeLoop(); err != nil {
st.l.Info("stream closed", zap.Error(err))
}
}()
return st, nil return st, nil
} }