From be5eadbeccf2904b2af21f2049d473ea967bddf1 Mon Sep 17 00:00:00 2001
From: Maisem Ali <maisem@tailscale.com>
Date: Tue, 2 Aug 2022 20:19:12 -0700
Subject: [PATCH] tsnet: log out ephemeral nodes on Close()

Signed-off-by: Maisem Ali <maisem@tailscale.com>
---
 tsnet/tsnet.go | 25 ++++++++++++++++++++-----
 1 file changed, 20 insertions(+), 5 deletions(-)

diff --git a/tsnet/tsnet.go b/tsnet/tsnet.go
index 0f8f1a03d..bd77e104e 100644
--- a/tsnet/tsnet.go
+++ b/tsnet/tsnet.go
@@ -130,6 +130,25 @@ func (s *Server) Start() error {
 //
 // It must not be called before or concurrently with Start.
 func (s *Server) Close() error {
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+	var wg sync.WaitGroup
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		// Perform a best-effort final flush.
+		s.logtail.Shutdown(ctx)
+	}()
+
+	if _, isMemStore := s.Store.(*mem.Store); isMemStore && s.Ephemeral {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			// Perform a best-effort logout.
+			s.lb.LogoutSync(ctx)
+		}()
+	}
+
 	s.shutdownCancel()
 	s.lb.Shutdown()
 	s.linkMon.Close()
@@ -143,11 +162,7 @@ func (s *Server) Close() error {
 	}
 	s.listeners = nil
 
-	// Perform a best-effort final flush.
-	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
-	defer cancel()
-	s.logtail.Shutdown(ctx)
-
+	wg.Wait()
 	return nil
 }