Support graceful terminator shutdown. Add more Xt ATs. Fixes GH-149 (#152)

* Support graceful terminator shutdown. Add more Xt ATs. Fixes GH-149

gateway/xgress_edge/accept.go

@@ -37,7 +37,7 @@ type edgeBindHandler struct {
 
 func (handler edgeBindHandler) BindChannel(ch channel2.Channel) error {
 	log := pfxlog.Logger()
-	log.WithField("token", ch.Id()).Info("accepting edge connection")
+	log.WithField("token", ch.Id()).Debug("accepting edge connection")
 
 	fpg := cert.NewFingerprintGenerator()
 	proxy := &ingressProxy{
@@ -64,6 +64,11 @@ func (handler edgeBindHandler) BindChannel(ch channel2.Channel) error {
 		Handler: proxy.processUnbind,
 	})
 
+	ch.AddReceiveHandler(&edge.AsyncFunctionReceiveAdapter{
+		Type:    edge.ContentTypeUpdateBind,
+		Handler: proxy.processUpdateBind,
+	})
+
 	ch.AddReceiveHandler(&edge.FunctionReceiveAdapter{
 		Type:    edge.ContentTypeStateClosed,
 		Handler: proxy.msgMux.HandleReceive,

gateway/xgress_edge/fabric.go

@@ -110,7 +110,7 @@ func (conn *localMessageSink) close(notify bool, reason string) {
 	}
 
 	log := pfxlog.ContextLogger(conn.Channel.Label()).WithField("connId", conn.Id())
-	log.Infof("closing message sink, reason: %v", reason)
+	log.Debugf("closing message sink, reason: %v", reason)
 	if notify {
 		// Notify edge client of close
 		closeMsg := edge.NewStateClosedMsg(conn.Id(), "")

gateway/xgress_edge/listener.go

@@ -296,6 +296,59 @@ func (proxy *ingressProxy) processUnbind(req *channel2.Message, ch channel2.Chan
 	}
 }
 
+func (proxy *ingressProxy) processUpdateBind(req *channel2.Message, ch channel2.Channel) {
+	token := string(req.Body)
+	log := pfxlog.ContextLogger(ch.Label()).WithField("sessionId", token).WithFields(edge.GetLoggerFields(req))
+
+	localListener, ok := proxy.listener.factory.hostedServices.Get(token)
+
+	if !ok {
+		log.Error("failed to update bind, no listener found")
+		return
+	}
+
+	sm := fabric.GetStateManager()
+	ns := sm.GetNetworkSession(token)
+
+	if ns == nil {
+		log.WithField("token", token).Error("session not found")
+		proxy.sendStateClosedReply("Invalid Session", req)
+		return
+	}
+
+	if _, found := proxy.fingerprints.HasAny(ns.CertFingerprints); !found {
+		log.WithField("token", token).
+			WithField("serviceFingerprints", ns.CertFingerprints).
+			WithField("clientFingerprints", proxy.fingerprints.Prints()).
+			Error("matching fingerprint not found")
+		proxy.sendStateClosedReply("Invalid Session", req)
+		return
+	}
+
+	var cost *uint16
+	if costBytes, hasCost := req.Headers[edge.CostHeader]; hasCost {
+		updatedCost := binary.LittleEndian.Uint16(costBytes)
+		cost = &updatedCost
+	}
+
+	var precedence *ctrl_pb.TerminatorPrecedence
+	if precedenceData, hasPrecedence := req.Headers[edge.PrecedenceHeader]; hasPrecedence && len(precedenceData) > 0 {
+		edgePrecedence := precedenceData[0]
+		updatedPrecedence := ctrl_pb.TerminatorPrecedence_Default
+		if edgePrecedence == edge.PrecedenceRequired {
+			updatedPrecedence = ctrl_pb.TerminatorPrecedence_Required
+		} else if edgePrecedence == edge.PrecedenceFailed {
+			updatedPrecedence = ctrl_pb.TerminatorPrecedence_Failed
+		}
+		precedence = &updatedPrecedence
+	}
+
+	log.Debugf("updating terminator %v to precedence %v and cost %v", localListener.terminatorId, precedence, cost)
+	if err := xgress.UpdateTerminator(proxy.listener.factory, localListener.terminatorId, cost, precedence); err != nil {
+		log.WithError(err).Error("failed to update bind")
+	}
+}
+
 func (proxy *ingressProxy) sendStateConnectedReply(req *channel2.Message, hostData map[uint32][]byte) {
 	connId, _ := req.GetUint32Header(edge.ConnIdHeader)
 	msg := edge.NewStateConnectedMsg(connId)

tests/context.go

@@ -36,6 +36,7 @@ import (
 	"github.com/netfoundry/ziti-foundation/util/concurrenz"
 	nfpem "github.com/netfoundry/ziti-foundation/util/pem"
 	sdkconfig "github.com/netfoundry/ziti-sdk-golang/ziti/config"
+	"github.com/netfoundry/ziti-sdk-golang/ziti/edge"
 	sdkenroll "github.com/netfoundry/ziti-sdk-golang/ziti/enroll"
 	"github.com/pkg/errors"
 	"io"
@@ -652,7 +653,7 @@ func (conn *testConn) RequireClose() {
 var testServerCounter uint64
 var testConnectionCounter uint64
 
-func newTestServer(listener net.Listener, dispatcher func(conn *testServerConn) error) *testServer {
+func newTestServer(listener edge.Listener, dispatcher func(conn *testServerConn) error) *testServer {
 	idx := atomic.AddUint64(&testServerCounter, 1)
 	return &testServer{
 		idx:        idx,
@@ -666,7 +667,7 @@ func newTestServer(listener net.Listener, dispatcher func(conn *testServerConn)
 
 type testServer struct {
 	idx        uint64
-	listener   net.Listener
+	listener   edge.Listener
 	errorC     chan error
 	closed     concurrenz.AtomicBoolean
 	msgCount   uint32
@@ -724,42 +725,42 @@ func (server *testServer) acceptLoop() {
 	select {
 	case _, ok := <-waitDone:
 		if !ok {
-			fmt.Println("all connections closed")
+			pfxlog.Logger().Debugf("all connections closed")
 		}
 	case <-time.After(10 * time.Second):
-		fmt.Println("timed out waiting for all connections to close")
+		pfxlog.Logger().Warn("timed out waiting for all connections to close")
 	}
 
 	close(server.errorC)
-	fmt.Printf("%v: service exiting\n", server.idx)
+	pfxlog.Logger().Debugf("%v: service exiting", server.idx)
 }
 
 func (server *testServer) dispatch(conn *testServerConn) {
 	defer func() {
-		fmt.Printf("marking waiter done to conn %v-%v\n", conn.server.idx, conn.id)
 		server.waiter.Done()
 	}()
 
+	log := pfxlog.Logger()
+
 	defer func() {
 		val := recover()
 		if val != nil {
 			if err, ok := val.(error); ok {
-				fmt.Printf("panic from server.dispatch: %+v\n", err)
+				log.WithError(err).Error("panic from server.dispatch")
 				server.errorC <- err
 			}
 		}
 	}()
 
 	defer func() {
-		fmt.Printf("closing conn %v-%v\n", conn.server.idx, conn.id)
 		conn.RequireClose()
 	}()
 
-	fmt.Printf("beginnging dispatch to conn %v-%v\n", conn.server.idx, conn.id)
+	log.Debugf("beginnging dispatch to conn %v-%v", conn.server.idx, conn.id)
 	err := server.dispatcher(conn)
-	fmt.Printf("finished dispatch to conn %v-%v\n", conn.server.idx, conn.id)
+	log.Debugf("finished dispatch to conn %v-%v", conn.server.idx, conn.id)
 	if err != nil {
-		fmt.Printf("failure from server.dispatch: %+v\n", err)
+		log.WithError(err).Error("failure from server.dispatch")
 		server.errorC <- err
 	}
 }

tests/data_flow_hs_rotating_test.go

@@ -0,0 +1,159 @@
+// +build apitests
+
+/*
+	Copyright NetFoundry, Inc.
+
+	Licensed under the Apache License, Version 2.0 (the "License");
+	you may not use this file except in compliance with the License.
+	You may obtain a copy of the License at
+
+	https://www.apache.org/licenses/LICENSE-2.0
+
+	Unless required by applicable law or agreed to in writing, software
+	distributed under the License is distributed on an "AS IS" BASIS,
+	WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+	See the License for the specific language governing permissions and
+	limitations under the License.
+*/
+
+package tests
+
+import (
+	"fmt"
+	"github.com/google/uuid"
+	"github.com/michaelquigley/pfxlog"
+	"github.com/netfoundry/ziti-foundation/util/concurrenz"
+	"github.com/netfoundry/ziti-sdk-golang/ziti"
+	"github.com/netfoundry/ziti-sdk-golang/ziti/edge"
+	"math/rand"
+	"sync/atomic"
+	"testing"
+	"time"
+)
+
+func Test_HSRotatingDataflow(t *testing.T) {
+	t.SkipNow()
+	ctx := NewTestContext(t)
+	defer ctx.teardown()
+	ctx.startServer()
+	ctx.requireAdminLogin()
+
+	ctx.createEnrollAndStartEdgeRouter()
+
+	t.Run("test smart routing", func(t *testing.T) {
+		ctx.testContextChanged(t)
+		testWithStrategy(ctx, "smartrouting")
+	})
+
+	t.Run("test random", func(t *testing.T) {
+		ctx.testContextChanged(t)
+		testWithStrategy(ctx, "random")
+	})
+
+	t.Run("test weighted", func(t *testing.T) {
+		ctx.testContextChanged(t)
+		testWithStrategy(ctx, "weighted")
+	})
+}
+
+func testWithStrategy(ctx *TestContext, strategy string) {
+	service := ctx.AdminSession.requireNewServiceAccessibleToAll(strategy)
+	fmt.Printf("service id: %v\n", service.id)
+
+	serverContextC := make(chan ziti.Context, 3)
+	doneC := make(chan struct{}, 1)
+	var serverContexts []ziti.Context
+	for i := 0; i < 3; i++ {
+		_, context := ctx.AdminSession.requireCreateSdkContext()
+		serverContexts = append(serverContexts, context)
+		serverContextC <- context
+	}
+
+	go func() {
+		logger := pfxlog.Logger()
+		for {
+			select {
+			case context := <-serverContextC:
+				listener, err := context.Listen(service.name)
+				ctx.req.NoError(err)
+				service := &rotatingService{
+					maxRequests: uint32(rand.Intn(5) + 5),
+					closeCB:     func() { serverContextC <- context },
+				}
+				server := newTestServer(listener, service.Handle)
+				server.start()
+				logger.Infof("started new listener, servicing %v reads", service.maxRequests)
+			case <-doneC:
+				break
+			}
+		}
+	}()
+
+	clientIdentity := ctx.AdminSession.requireNewIdentityWithOtt(false)
+	clientConfig := ctx.enrollIdentity(clientIdentity.id)
+	clientContext := ziti.NewContextWithConfig(clientConfig)
+
+	for i := 0; i < 250; i++ {
+		// fmt.Printf("%v: start\n", i+1)
+		conn := ctx.wrapConn(clientContext.Dial(service.name))
+
+		name := uuid.New().String()
+		conn.WriteString(name, time.Second)
+		conn.ReadExpected("hello, "+name, time.Second)
+		conn.RequireClose()
+		fmt.Printf("%v: done\n", i+1)
+	}
+
+	close(doneC)
+
+	for range serverContexts {
+		conn := ctx.wrapConn(clientContext.Dial(service.name))
+		conn.WriteString("quit", time.Second)
+		conn.ReadExpected("ok", time.Second)
+	}
+}
+
+type rotatingService struct {
+	maxRequests uint32
+	requests    uint32
+	closeCB     func()
+	closing     concurrenz.AtomicBoolean
+}
+
+func (service *rotatingService) Handle(conn *testServerConn) error {
+	requests := atomic.AddUint32(&service.requests, 1)
+	doClose := requests >= service.maxRequests
+	logger := pfxlog.Logger()
+	for {
+		name, eof := conn.ReadString(1024, time.Minute)
+		if eof {
+			return nil
+		}
+
+		if name == "quit" {
+			logger.Infof("%v-%v: received '%v' from client", conn.server.idx, conn.id, name)
+			if err := conn.server.listener.UpdatePrecedence(edge.PrecedenceFailed); err != nil {
+				logger.WithError(err).Error("failed to update precedence")
+				return err
+			}
+			conn.WriteString("ok", time.Second)
+			return conn.server.close()
+		}
+
+		conn.WriteString("hello, "+name, time.Second)
+		atomic.AddUint32(&conn.server.msgCount, 1)
+		if doClose && service.closing.CompareAndSwap(false, true) {
+			logger.Infof("%v-%v: preparing to exit after serving %v requests. setting precedence to failed",
+				conn.server.idx, conn.id, atomic.LoadUint32(&service.requests))
+			if err := conn.server.listener.UpdatePrecedence(edge.PrecedenceFailed); err != nil {
+				logger.WithError(err).Error("failed to update precedence")
+				return err
+			}
+			service.closeCB()
+			logger.Debugf("%v-%v: sleeping", conn.server.idx, conn.id)
+			time.Sleep(100 * time.Millisecond)
+			logger.Debugf("%v-%v: exiting", conn.server.idx, conn.id)
+			return conn.server.close()
+		}
+	}
+}