Reverse: portal-worker should not be closed until there is at least one other worker

This commit is contained in:
patterniha 2025-07-07 21:57:47 +03:30
parent cb1afb33e6
commit 8844058e27

View file

@ -76,7 +76,7 @@ func (p *Portal) HandleConnection(ctx context.Context, link *transport.Link) err
return errors.New("failed to create mux client worker").Base(err).AtWarning() return errors.New("failed to create mux client worker").Base(err).AtWarning()
} }
worker, err := NewPortalWorker(muxClient) worker, err := NewPortalWorker(muxClient, p.picker)
if err != nil { if err != nil {
return errors.New("failed to create portal worker").Base(err) return errors.New("failed to create portal worker").Base(err)
} }
@ -170,7 +170,7 @@ func (p *StaticMuxPicker) PickAvailable() (*mux.ClientWorker, error) {
if w.draining { if w.draining {
continue continue
} }
if w.client.Closed() { if w.IsFull() {
continue continue
} }
if w.client.ActiveConnections() < minConn { if w.client.ActiveConnections() < minConn {
@ -206,14 +206,16 @@ func (p *StaticMuxPicker) AddWorker(worker *PortalWorker) {
} }
type PortalWorker struct { type PortalWorker struct {
client *mux.ClientWorker client *mux.ClientWorker
control *task.Periodic control *task.Periodic
writer buf.Writer writer buf.Writer
reader buf.Reader reader buf.Reader
draining bool draining bool
counter uint32
parentPicker *StaticMuxPicker
} }
func NewPortalWorker(client *mux.ClientWorker) (*PortalWorker, error) { func NewPortalWorker(client *mux.ClientWorker, picker *StaticMuxPicker) (*PortalWorker, error) {
opt := []pipe.Option{pipe.WithSizeLimit(16 * 1024)} opt := []pipe.Option{pipe.WithSizeLimit(16 * 1024)}
uplinkReader, uplinkWriter := pipe.New(opt...) uplinkReader, uplinkWriter := pipe.New(opt...)
downlinkReader, downlinkWriter := pipe.New(opt...) downlinkReader, downlinkWriter := pipe.New(opt...)
@ -231,9 +233,10 @@ func NewPortalWorker(client *mux.ClientWorker) (*PortalWorker, error) {
return nil, errors.New("unable to dispatch control connection") return nil, errors.New("unable to dispatch control connection")
} }
w := &PortalWorker{ w := &PortalWorker{
client: client, client: client,
reader: downlinkReader, reader: downlinkReader,
writer: uplinkWriter, writer: uplinkWriter,
parentPicker: picker,
} }
w.control = &task.Periodic{ w.control = &task.Periodic{
Execute: w.heartbeat, Execute: w.heartbeat,
@ -244,32 +247,32 @@ func NewPortalWorker(client *mux.ClientWorker) (*PortalWorker, error) {
} }
func (w *PortalWorker) heartbeat() error { func (w *PortalWorker) heartbeat() error {
if w.client.Closed() { if w.Closed() {
return errors.New("client worker stopped") return errors.New("client worker stopped")
} }
if w.draining || w.writer == nil { if w.writer == nil {
return errors.New("already disposed") return errors.New("already disposed")
} }
msg := &Control{} msg := &Control{}
msg.FillInRandom() msg.FillInRandom()
if w.client.TotalConnections() > 256 { if w.draining || w.client.TotalConnections() > 256 {
w.draining = true w.draining = true
msg.State = Control_DRAIN msg.State = Control_DRAIN
defer func() { defer w.tryCloseHeartbeat()
common.Close(w.writer)
common.Interrupt(w.reader)
w.writer = nil
}()
} }
b, err := proto.Marshal(msg) w.counter = (w.counter + 1) % 5
common.Must(err) if w.draining || w.counter == 1 {
mb := buf.MergeBytes(nil, b) b, err := proto.Marshal(msg)
return w.writer.WriteMultiBuffer(mb) common.Must(err)
mb := buf.MergeBytes(nil, b)
return w.writer.WriteMultiBuffer(mb)
}
return nil
} }
func (w *PortalWorker) IsFull() bool { func (w *PortalWorker) IsFull() bool {
@ -279,3 +282,20 @@ func (w *PortalWorker) IsFull() bool {
func (w *PortalWorker) Closed() bool { func (w *PortalWorker) Closed() bool {
return w.client.Closed() return w.client.Closed()
} }
func (w *PortalWorker) tryCloseHeartbeat() {
w.parentPicker.access.Lock()
closeable := false
for _, wo := range w.parentPicker.workers {
if wo != w && !wo.IsFull() && !wo.draining && wo.writer != nil {
closeable = true
break
}
}
w.parentPicker.access.Unlock()
if closeable {
common.Close(w.writer)
common.Interrupt(w.reader)
w.writer = nil
}
}