From b93513a9b4fd360749f5022965874769db21dc84 Mon Sep 17 00:00:00 2001
From: Neil Alexander <neilalexander@users.noreply.github.com>
Date: Wed, 5 Aug 2020 16:26:50 +0100
Subject: [PATCH] Try to wake up queues on incoming QUIC session

---
 cmd/dendrite-demo-yggdrasil/main.go            | 12 ++++++++++++
 cmd/dendrite-demo-yggdrasil/yggconn/node.go    |  7 ++++++-
 cmd/dendrite-demo-yggdrasil/yggconn/session.go |  2 ++
 3 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/cmd/dendrite-demo-yggdrasil/main.go b/cmd/dendrite-demo-yggdrasil/main.go
index 78ecc9078..73bfceff6 100644
--- a/cmd/dendrite-demo-yggdrasil/main.go
+++ b/cmd/dendrite-demo-yggdrasil/main.go
@@ -122,6 +122,18 @@ func main() {
 		base, federation, rsAPI, stateAPI, keyRing,
 	)
 
+	ygg.SetSessionFunc(func(address string) {
+		req := &api.PerformServersAliveRequest{
+			Servers: []gomatrixserverlib.ServerName{
+				gomatrixserverlib.ServerName(address),
+			},
+		}
+		res := &api.PerformServersAliveResponse{}
+		if err := fsAPI.PerformServersAlive(context.TODO(), req, res); err != nil {
+			logrus.WithError(err).Error("Failed to send wake-up message to newly connected node")
+		}
+	})
+
 	rsComponent.SetFederationSenderAPI(fsAPI)
 
 	embed.Embed(base.BaseMux, *instancePort, "Yggdrasil Demo")
diff --git a/cmd/dendrite-demo-yggdrasil/yggconn/node.go b/cmd/dendrite-demo-yggdrasil/yggconn/node.go
index 697813c3c..6949d4d87 100644
--- a/cmd/dendrite-demo-yggdrasil/yggconn/node.go
+++ b/cmd/dendrite-demo-yggdrasil/yggconn/node.go
@@ -52,6 +52,7 @@ type Node struct {
 	quicConfig   *quic.Config
 	sessions     sync.Map // string -> quic.Session
 	sessionCount atomic.Uint32
+	sessionFunc  func(address string)
 	coords       sync.Map // string -> yggdrasil.Coords
 	incoming     chan QUICStream
 	NewSession   func(remote gomatrixserverlib.ServerName)
@@ -176,6 +177,10 @@ func (n *Node) SigningPrivateKey() ed25519.PrivateKey {
 	return ed25519.PrivateKey(privBytes)
 }
 
+func (n *Node) SetSessionFunc(f func(address string)) {
+	n.sessionFunc = f
+}
+
 func (n *Node) PeerCount() int {
 	return len(n.core.GetPeers()) - 1
 }
@@ -186,7 +191,7 @@ func (n *Node) SessionCount() int {
 
 func (n *Node) KnownNodes() []gomatrixserverlib.ServerName {
 	nodemap := map[string]struct{}{
-		"b5ae50589e50991dd9dd7d59c5c5f7a4521e8da5b603b7f57076272abc58b374": struct{}{},
+		"b5ae50589e50991dd9dd7d59c5c5f7a4521e8da5b603b7f57076272abc58b374": {},
 	}
 	/*
 		for _, peer := range n.core.GetSwitchPeers() {
diff --git a/cmd/dendrite-demo-yggdrasil/yggconn/session.go b/cmd/dendrite-demo-yggdrasil/yggconn/session.go
index 6c7b427d5..22b5646bf 100644
--- a/cmd/dendrite-demo-yggdrasil/yggconn/session.go
+++ b/cmd/dendrite-demo-yggdrasil/yggconn/session.go
@@ -59,6 +59,7 @@ func (n *Node) listenFromYgg() {
 		address := session.ConnectionState().PeerCertificates[0].Subject.CommonName
 		n.log.Infoln("Accepted connection from", address)
 		go n.listenFromQUIC(session, address)
+		go n.sessionFunc(address)
 	}
 }
 
@@ -209,6 +210,7 @@ func (n *Node) tryDial(address string, coords yggdrasil.Coords) (quic.Session, e
 		return nil, fmt.Errorf("expected %q but dialled %q", address, gotAddress)
 	}
 	go n.listenFromQUIC(session, address)
+	go n.sessionFunc(address)
 	return session, nil
 }