From 8d3ea033063d653148472f00486248eac1f76afe Mon Sep 17 00:00:00 2001
From: Paul Lorenz <paul.lorenz@netfoundry.io>
Date: Tue, 9 Mar 2021 12:10:04 -0500
Subject: [PATCH] Use bitset for xgress flags

---
 router/xgress/link_send_buffer.go |  2 +-
 router/xgress/retransmitter.go    |  2 +-
 router/xgress/xgress.go           | 50 +++++++++++--------------------
 3 files changed, 19 insertions(+), 35 deletions(-)

diff --git a/router/xgress/link_send_buffer.go b/router/xgress/link_send_buffer.go
index 0e96ac69c..ee7197bb1 100644
--- a/router/xgress/link_send_buffer.go
+++ b/router/xgress/link_send_buffer.go
@@ -207,7 +207,7 @@ func (buffer *LinkSendBuffer) run() {
 		case ack := <-buffer.newlyReceivedAcks:
 			buffer.receiveAcknowledgement(ack)
 			buffer.retransmit()
-			if buffer.closeWhenEmpty.Get() && len(buffer.buffer) == 0 && !buffer.x.closed.Get() && buffer.x.IsEndOfSessionSent() {
+			if buffer.closeWhenEmpty.Get() && len(buffer.buffer) == 0 && !buffer.x.Closed() && buffer.x.IsEndOfSessionSent() {
 				go buffer.x.Close()
 			}
 
diff --git a/router/xgress/retransmitter.go b/router/xgress/retransmitter.go
index 5cafa8052..2d1a0e6a3 100644
--- a/router/xgress/retransmitter.go
+++ b/router/xgress/retransmitter.go
@@ -147,7 +147,7 @@ func (retransmitter *Retransmitter) retransmitSender() {
 			if !retransmit.isAcked() {
 				if err := retransmitter.forwarder.ForwardPayload(retransmit.x.address, retransmit.payload); err != nil {
 					// if xgress is closed, don't log the error. We still want to try retransmitting in case we're re-sending end of session
-					if !retransmit.x.closed.Get() {
+					if !retransmit.x.Closed() {
 						logger.WithError(err).Errorf("unexpected error while retransmitting payload from [@/%v]", retransmit.x.address)
 						retransmissionFailures.Mark(1)
 						retransmitter.faultReporter.ReportForwardingFault(retransmit.payload.SessionId)
diff --git a/router/xgress/xgress.go b/router/xgress/xgress.go
index f60541eb9..dac879069 100644
--- a/router/xgress/xgress.go
+++ b/router/xgress/xgress.go
@@ -34,6 +34,11 @@ import (
 
 const (
 	HeaderKeyUUID = 0
+
+	closedFlag            = 0
+	rxerStartedFlag       = 1
+	endOfSessionRecvdFlag = 2
+	endOfSessionSentFlag  = 3
 )
 
 type Address string
@@ -107,11 +112,7 @@ type Xgress struct {
 	linkRxBuffer         *LinkReceiveBuffer
 	closeHandler         CloseHandler
 	peekHandlers         []PeekHandler
-	closed               concurrenz.AtomicBoolean
-	rxerStarted          bool
-	endOfSessionRecvd    bool
-	endOfSessionSent     bool
-	flagsMutex           sync.Mutex
+	flags                concurrenz.AtomicBitSet
 	timeOfLastRxFromLink int64
 }
 
@@ -165,31 +166,19 @@ func (self *Xgress) AddPeekHandler(peekHandler PeekHandler) {
 }
 
 func (self *Xgress) IsEndOfSessionReceived() bool {
-	self.flagsMutex.Lock()
-	defer self.flagsMutex.Unlock()
-	return self.endOfSessionRecvd
+	return self.flags.IsSet(endOfSessionRecvdFlag)
 }
 
 func (self *Xgress) markSessionEndReceived() {
-	self.flagsMutex.Lock()
-	defer self.flagsMutex.Unlock()
-	self.endOfSessionRecvd = true
+	self.flags.Set(endOfSessionRecvdFlag, true)
 }
 
 func (self *Xgress) IsSessionStarted() bool {
-	self.flagsMutex.Lock()
-	defer self.flagsMutex.Unlock()
-	return !self.IsTerminator() || self.rxerStarted
+	return !self.IsTerminator() || self.flags.IsSet(rxerStartedFlag)
 }
 
 func (self *Xgress) firstSessionStartReceived() bool {
-	self.flagsMutex.Lock()
-	defer self.flagsMutex.Unlock()
-	if self.rxerStarted {
-		return false
-	}
-	self.rxerStarted = true
-	return true
+	return self.flags.CompareAndSet(rxerStartedFlag, false, true)
 }
 
 func (self *Xgress) Start() {
@@ -243,20 +232,15 @@ func (self *Xgress) GetEndSession() *Payload {
 }
 
 func (self *Xgress) ForwardEndOfSession(sendF func(payload *Payload) bool) {
-	self.flagsMutex.Lock()
-	defer self.flagsMutex.Unlock()
-
 	// for now always send end of session. too many is better than not enough
-	if !self.endOfSessionRecvd {
+	if !self.IsEndOfSessionSent() {
 		sendF(self.GetEndSession())
-		self.endOfSessionSent = true
+		self.flags.Set(endOfSessionSentFlag, true)
 	}
 }
 
 func (self *Xgress) IsEndOfSessionSent() bool {
-	self.flagsMutex.Lock()
-	defer self.flagsMutex.Unlock()
-	return self.endOfSessionSent
+	return self.flags.IsSet(endOfSessionSentFlag)
 }
 
 func (self *Xgress) CloseTimeout(duration time.Duration) {
@@ -277,7 +261,7 @@ Things which can trigger close
 func (self *Xgress) Close() {
 	log := pfxlog.ContextLogger(self.Label())
 
-	if self.closed.CompareAndSwap(false, true) {
+	if self.flags.CompareAndSet(closedFlag, false, true) {
 		log.Debug("closing xgress peer")
 		if err := self.peer.Close(); err != nil {
 			log.WithError(err).Warn("error while closing xgress peer")
@@ -301,11 +285,11 @@ func (self *Xgress) Close() {
 }
 
 func (self *Xgress) Closed() bool {
-	return self.closed.Get()
+	return self.flags.IsSet(closedFlag)
 }
 
 func (self *Xgress) SendPayload(payload *Payload) error {
-	if self.closed.Get() {
+	if self.Closed() {
 		return nil
 	}
 
@@ -479,7 +463,7 @@ func (self *Xgress) rx() {
 			return
 		}
 
-		if self.closed.Get() {
+		if self.Closed() {
 			return
 		}
 		start := 0