From 2238a288d9b4b43417a4c6e9673344b8490d7180 Mon Sep 17 00:00:00 2001
From: Jakob Borg <jakob@kastelo.net>
Date: Sun, 15 Sep 2024 11:37:49 +0200
Subject: [PATCH] fix(model): shut down index sender faster (#9704)

---
 lib/model/indexhandler.go |  6 ++++++
 lib/protocol/protocol.go  | 12 ++++++++++++
 2 files changed, 18 insertions(+)

diff --git a/lib/model/indexhandler.go b/lib/model/indexhandler.go
index 00d92eb6f..55a55b2ed 100644
--- a/lib/model/indexhandler.go
+++ b/lib/model/indexhandler.go
@@ -233,6 +233,12 @@ func (s *indexHandler) sendIndexTo(ctx context.Context, fset *db.FileSet) error
 	batch := db.NewFileInfoBatch(nil)
 	var batchError error
 	batch.SetFlushFunc(func(fs []protocol.FileInfo) error {
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		default:
+		}
+
 		if len(fs) == 0 {
 			// can't happen, flush is not called with an empty batch
 			panic("bug: flush called with empty batch (race condition?)")
diff --git a/lib/protocol/protocol.go b/lib/protocol/protocol.go
index f7f725f1e..b8e399fe1 100644
--- a/lib/protocol/protocol.go
+++ b/lib/protocol/protocol.go
@@ -354,6 +354,8 @@ func (c *rawConnection) Index(ctx context.Context, idx *Index) error {
 	select {
 	case <-c.closed:
 		return ErrClosed
+	case <-ctx.Done():
+		return ctx.Err()
 	default:
 	}
 	c.idxMut.Lock()
@@ -367,6 +369,8 @@ func (c *rawConnection) IndexUpdate(ctx context.Context, idxUp *IndexUpdate) err
 	select {
 	case <-c.closed:
 		return ErrClosed
+	case <-ctx.Done():
+		return ctx.Err()
 	default:
 	}
 	c.idxMut.Lock()
@@ -377,6 +381,14 @@ func (c *rawConnection) IndexUpdate(ctx context.Context, idxUp *IndexUpdate) err
 
 // Request returns the bytes for the specified block after fetching them from the connected peer.
 func (c *rawConnection) Request(ctx context.Context, req *Request) ([]byte, error) {
+	select {
+	case <-c.closed:
+		return nil, ErrClosed
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	default:
+	}
+
 	rc := make(chan asyncResult, 1)
 
 	c.awaitingMut.Lock()