Merge pull request openziti/fabric#144 from openziti/xgress-delayed-close

Add delayed close handling for xgress

router/xgress/link_send_buffer.go

@@ -17,12 +17,10 @@
 package xgress
 
 import (
-	"fmt"
 	"github.com/michaelquigley/pfxlog"
 	"github.com/openziti/foundation/util/concurrenz"
 	"github.com/openziti/foundation/util/info"
 	"github.com/pkg/errors"
-	"github.com/sirupsen/logrus"
 	"math"
 	"sync/atomic"
 	"time"
@@ -52,6 +50,7 @@ type LinkSendBuffer struct {
 	retxScale             float64
 	retxThreshold         uint32
 	lastRetransmitTime    int64
+	closeWhenEmpty        concurrenz.AtomicBoolean
 }
 
 type txPayload struct {
@@ -108,11 +107,15 @@ func NewLinkSendBuffer(x *Xgress) *LinkSendBuffer {
 	return buffer
 }
 
+func (buffer *LinkSendBuffer) CloseWhenEmpty() {
+	buffer.closeWhenEmpty.Set(true)
+}
+
 func (buffer *LinkSendBuffer) BufferPayload(payload *Payload) (func(), error) {
 	txPayload := &txPayload{payload: payload, age: math.MaxInt64, x: buffer.x}
 	select {
 	case buffer.newlyBuffered <- txPayload:
-		pfxlog.ContextLogger("s/"+payload.GetSessionId()).Debugf("buffered [%d]", payload.GetSequence())
+		pfxlog.ContextLogger(buffer.x.Label()).Debugf("buffered [%d]", payload.GetSequence())
 		return txPayload.markSent, nil
 	case <-buffer.closeNotify:
 		return nil, errors.Errorf("payload buffer closed")
@@ -120,7 +123,7 @@ func (buffer *LinkSendBuffer) BufferPayload(payload *Payload) (func(), error) {
 }
 
 func (buffer *LinkSendBuffer) ReceiveAcknowledgement(ack *Acknowledgement) {
-	log := pfxlog.Logger().WithFields(ack.GetLoggerFields())
+	log := pfxlog.ContextLogger(buffer.x.Label()).WithFields(ack.GetLoggerFields())
 	log.Debug("ack received")
 	select {
 	case buffer.newlyReceivedAcks <- ack:
@@ -131,7 +134,7 @@ func (buffer *LinkSendBuffer) ReceiveAcknowledgement(ack *Acknowledgement) {
 }
 
 func (buffer *LinkSendBuffer) Close() {
-	logrus.Debugf("[%p] closing", buffer)
+	pfxlog.ContextLogger(buffer.x.Label()).Debugf("[%p] closing", buffer)
 	if buffer.closed.CompareAndSwap(false, true) {
 		close(buffer.closeNotify)
 	}
@@ -163,14 +166,14 @@ func (buffer *LinkSendBuffer) isBlocked() bool {
 	}
 
 	if blocked {
-		pfxlog.Logger().Debugf("blocked=%v win_size=%v tx_buffer_size=%v rx_buffer_size=%v", blocked, buffer.windowsSize, buffer.linkRecvBufferSize, buffer.linkSendBufferSize)
+		pfxlog.ContextLogger(buffer.x.Label()).Debugf("blocked=%v win_size=%v tx_buffer_size=%v rx_buffer_size=%v", blocked, buffer.windowsSize, buffer.linkRecvBufferSize, buffer.linkSendBufferSize)
 	}
 
 	return blocked
 }
 
 func (buffer *LinkSendBuffer) run() {
-	log := pfxlog.ContextLogger("s/" + buffer.x.sessionId)
+	log := pfxlog.ContextLogger(buffer.x.Label())
 	defer log.Debugf("[%p] exited", buffer)
 	log.Debugf("[%p] started", buffer)
 
@@ -205,6 +208,9 @@ func (buffer *LinkSendBuffer) run() {
 		case ack := <-buffer.newlyReceivedAcks:
 			buffer.receiveAcknowledgement(ack)
 			buffer.retransmit()
+			if buffer.closeWhenEmpty.Get() && len(buffer.buffer) == 0 {
+				go buffer.x.Close()
+			}
 
 		case txPayload := <-buffered:
 			buffer.buffer[txPayload.payload.GetSequence()] = txPayload
@@ -235,7 +241,7 @@ func (buffer *LinkSendBuffer) close() {
 }
 
 func (buffer *LinkSendBuffer) receiveAcknowledgement(ack *Acknowledgement) {
-	log := pfxlog.Logger().WithFields(ack.GetLoggerFields())
+	log := pfxlog.ContextLogger(buffer.x.Label()).WithFields(ack.GetLoggerFields())
 
 	for _, sequence := range ack.Sequence {
 		if txPayload, found := buffer.buffer[sequence]; found {
@@ -289,7 +295,7 @@ func (buffer *LinkSendBuffer) receiveAcknowledgement(ack *Acknowledgement) {
 func (buffer *LinkSendBuffer) retransmit() {
 	now := info.NowInMilliseconds()
 	if len(buffer.buffer) > 0 && (now-buffer.lastRetransmitTime) > 64 {
-		log := pfxlog.ContextLogger(fmt.Sprintf("s/" + buffer.x.sessionId))
+		log := pfxlog.ContextLogger(buffer.x.Label())
 
 		retransmitted := 0
 		for _, v := range buffer.buffer {

router/xgress/options.go

@@ -18,6 +18,7 @@ package xgress
 
 import (
 	"encoding/json"
+	"time"
 )
 
 // Options contains common Xgress configuration options
@@ -41,6 +42,7 @@ type Options struct {
 	RetxInitial  uint32
 	RetxScale    float64
 	RetxAddMs    uint32
+	MaxCloseWait time.Duration
 }
 
 func LoadOptions(data OptionsData) *Options {
@@ -102,6 +104,10 @@ func LoadOptions(data OptionsData) *Options {
 		if value, found := data["retxAddMs"]; found {
 			options.RetxAddMs = uint32(value.(int))
 		}
+		if value, found := data["maxCloseWaitMs"]; found {
+			options.MaxCloseWait = time.Duration(value.(int)) * time.Millisecond
+		}
+
 	}
 
 	return options
@@ -125,6 +131,7 @@ func DefaultOptions() *Options {
 		RxBufferSize:           4 * 1024 * 1024,
 		RetxScale:              2.0,
 		RetxAddMs:              100,
+		MaxCloseWait:           30 * time.Second,
 	}
 }
 

router/xgress/xgress.go

@@ -311,20 +311,30 @@ func (self *Xgress) tx() {
 
 	log.Debug("started")
 	defer log.Debug("exited")
+	defer self.Close()
 
 	var payload *Payload
 
 	for {
 		payload = self.nextPayload()
 
-		if payload == nil || payload.IsSessionEndFlagSet() {
+		if payload == nil {
+			log.Debug("nil payload received, exiting")
+			return
+		}
+
+		if payload.IsSessionEndFlagSet() {
+			log.Debug("session end payload recieved, exiting")
 			return
 		}
 
 		payloadLogger := log.WithFields(payload.GetLoggerFields())
+		payloadLogger.Debug("sending")
+
 		for _, peekHandler := range self.peekHandlers {
 			peekHandler.Tx(self, payload)
 		}
+
 		if !payload.IsSessionStartFlagSet() && !payload.IsSessionEndFlagSet() {
 			start := time.Now()
 			n, err := self.peer.WritePayload(payload.Data, payload.Headers)
@@ -338,8 +348,8 @@ func (self *Xgress) tx() {
 			}
 		}
 		payloadSize := len(payload.Data)
-		size := atomic.AddUint32(&self.linkRxBuffer.size, ^uint32(payloadSize-1))
-		pfxlog.Logger().Debugf("Payload %v of size %v removed from transmit buffer. New size: %v", payload.Sequence, payloadSize, size)
+		size := atomic.AddUint32(&self.linkRxBuffer.size, ^uint32(payloadSize-1)) // subtraction for uint32
+		payloadLogger.Debugf("Payload %v of size %v removed from rx buffer. New size: %v", payload.Sequence, payloadSize, size)
 		localRecvBufferSizeBytesHistogram.Update(int64(size))
 
 		lastBufferSizeSent := self.linkRxBuffer.getLastBufferSizeSent()
@@ -361,7 +371,8 @@ func (self *Xgress) rx() {
 			return
 		}
 	}()
-	defer self.Close()
+	defer self.CloseTimeout(self.Options.MaxCloseWait)
+	defer func() { self.forwardPayload(self.GetEndSession()) }()
 
 	for {
 		buffer, headers, err := self.peer.ReadPayload()
@@ -417,7 +428,7 @@ func (self *Xgress) forwardPayload(payload *Payload) {
 	sendCallback, err := self.payloadBuffer.BufferPayload(payload)
 
 	if err != nil {
-		pfxlog.Logger().WithError(err).Error("failure forwarding payload")
+		pfxlog.ContextLogger(self.Label()).WithError(err).Error("failure forwarding payload")
 		return
 	}
 
@@ -440,15 +451,15 @@ func (self *Xgress) nextReceiveSequence() int32 {
 }
 
 func (self *Xgress) closeTimeoutHandler(duration time.Duration) {
+	self.payloadBuffer.CloseWhenEmpty() // If we clear the send buffer, close sooner
 	time.Sleep(duration)
 	self.Close()
 }
 
 func (self *Xgress) PayloadReceived(payload *Payload) {
-	log := pfxlog.Logger().WithFields(payload.GetLoggerFields())
+	log := pfxlog.ContextLogger(self.Label()).WithFields(payload.GetLoggerFields())
 	log.Debug("payload received")
 	if self.linkRxBuffer.ReceiveUnordered(payload, self.Options.RxBufferSize) {
-		log := pfxlog.Logger().WithFields(payload.GetLoggerFields())
 		log.Debug("ready to acknowledge")
 
 		ack := NewAcknowledgement(self.sessionId, self.originator)
@@ -458,11 +469,13 @@ func (self *Xgress) PayloadReceived(payload *Payload) {
 
 		atomic.StoreUint32(&self.linkRxBuffer.lastBufferSizeSent, ack.RecvBufferSize)
 		acker.ack(ack, self.address)
+	} else {
+		log.Debug("dropped")
 	}
 }
 
 func (self *Xgress) SendEmptyAck() {
-	pfxlog.Logger().WithField("session", self.sessionId).Debug("sending empty ack")
+	pfxlog.ContextLogger(self.Label()).WithField("session", self.sessionId).Debug("sending empty ack")
 	ack := NewAcknowledgement(self.sessionId, self.originator)
 	ack.RecvBufferSize = self.linkRxBuffer.Size()
 	atomic.StoreUint32(&self.linkRxBuffer.lastBufferSizeSent, ack.RecvBufferSize)