73 lines
1.5 KiB
Go
73 lines
1.5 KiB
Go
package streampool
|
|
|
|
import (
|
|
"context"
|
|
"sync/atomic"
|
|
|
|
"github.com/anyproto/any-sync/app/logger"
|
|
"github.com/cheggaaa/mb/v3"
|
|
"go.uber.org/zap"
|
|
"storj.io/drpc"
|
|
)
|
|
|
|
type stream struct {
|
|
peerId string
|
|
peerCtx context.Context
|
|
stream drpc.Stream
|
|
pool *streamPool
|
|
streamId uint32
|
|
closed atomic.Bool
|
|
l logger.CtxLogger
|
|
queue *mb.MB[drpc.Message]
|
|
tags []string
|
|
}
|
|
|
|
func (sr *stream) write(msg drpc.Message) (err error) {
|
|
return sr.queue.TryAdd(msg)
|
|
}
|
|
|
|
func (sr *stream) readLoop() error {
|
|
defer func() {
|
|
sr.streamClose()
|
|
}()
|
|
sr.l.Debug("stream read started")
|
|
for {
|
|
msg := sr.pool.handler.NewReadMessage()
|
|
if err := sr.stream.MsgRecv(msg, EncodingProto); err != nil {
|
|
sr.l.Info("msg receive error", zap.Error(err))
|
|
return err
|
|
}
|
|
ctx := streamCtx(sr.peerCtx, sr.streamId, sr.peerId)
|
|
ctx = logger.CtxWithFields(ctx, zap.String("peerId", sr.peerId))
|
|
if err := sr.pool.handler.HandleMessage(ctx, sr.peerId, msg); err != nil {
|
|
sr.l.Info("msg handle error", zap.Error(err))
|
|
return err
|
|
}
|
|
}
|
|
}
|
|
|
|
func (sr *stream) writeLoop() {
|
|
for {
|
|
msg, err := sr.queue.WaitOne(sr.peerCtx)
|
|
if err != nil {
|
|
if err != mb.ErrClosed {
|
|
sr.streamClose()
|
|
}
|
|
return
|
|
}
|
|
if err := sr.stream.MsgSend(msg, EncodingProto); err != nil {
|
|
sr.l.Warn("msg send error", zap.Error(err))
|
|
sr.streamClose()
|
|
return
|
|
}
|
|
}
|
|
}
|
|
|
|
func (sr *stream) streamClose() {
|
|
if !sr.closed.Swap(true) {
|
|
_ = sr.queue.Close()
|
|
_ = sr.stream.Close()
|
|
sr.pool.removeStream(sr.streamId)
|
|
}
|
|
}
|