From 69708b0e1ac425c937f20f1cd581c3cff757576c Mon Sep 17 00:00:00 2001 From: Zeyad Gouda Date: Mon, 6 Apr 2026 18:47:09 +0200 Subject: [PATCH 1/7] tests: add /debug endpoint to fakestore to allow interrupting downloads Signed-off-by: Zeyad Gouda --- tests/lib/fakestore/store/store.go | 120 +++++++++++++++++++++++- tests/lib/fakestore/store/store_test.go | 90 ++++++++++++++++++ 2 files changed, 209 insertions(+), 1 deletion(-) diff --git a/tests/lib/fakestore/store/store.go b/tests/lib/fakestore/store/store.go index 8d4cc707cdf..08c894feb19 100644 --- a/tests/lib/fakestore/store/store.go +++ b/tests/lib/fakestore/store/store.go @@ -79,6 +79,8 @@ type Store struct { channelRepository *ChannelRepository snapsCache map[string]snapCachedInfo + + killAfter map[string]int64 } // NewStore creates a new store server serving snaps from the given top directory and assertions from topDir/asserts. If assertFallback is true missing assertions are looked up in the main online store. @@ -105,13 +107,16 @@ func NewStore(topDir, addr string, assertFallback bool) *Store { rootDir: filepath.Join(topDir, "channels"), }, snapsCache: make(map[string]snapCachedInfo), + killAfter: make(map[string]int64), } mux.HandleFunc("/", rootEndpoint) mux.HandleFunc("/api/v1/snaps/search", store.searchEndpoint) mux.HandleFunc("/api/v1/snaps/details/", store.detailsEndpoint) mux.HandleFunc("/api/v1/snaps/metadata", store.bulkEndpoint) - mux.Handle("/download/", http.StripPrefix("/download/", http.FileServer(http.Dir(topDir)))) + + fileServer := http.StripPrefix("/download/", http.FileServer(http.Dir(topDir))) + mux.Handle("/download/", logRangeHeader(store.applyKillAfter(fileServer.ServeHTTP))) mux.HandleFunc("/api/v1/snaps/auth/nonces", store.nonceEndpoint) mux.HandleFunc("/api/v1/snaps/auth/sessions", store.sessionEndpoint) @@ -122,6 +127,8 @@ func NewStore(topDir, addr string, assertFallback bool) *Store { mux.HandleFunc("/v2/repairs/", store.repairsEndpoint) + mux.HandleFunc("/debug", store.debugEndpoint) + return store } @@ -354,6 +361,117 @@ type detailsReplyJSON struct { Base string `json:"base,omitempty"` } +type killAfterWriter struct { + http.ResponseWriter + killAfter int64 +} + +func (kaw *killAfterWriter) Write(p []byte) (int, error) { + if kaw.killAfter >= 0 { + kaw.killAfter -= int64(len(p)) + } + + if kaw.killAfter < 0 { + // hijack the connection to force a hard drop + hj, ok := kaw.ResponseWriter.(http.Hijacker) + if ok { + conn, _, _ := hj.Hijack() + conn.Close() // hard close the TCP connection + } + return 0, fmt.Errorf("connection killed") + } + + return kaw.ResponseWriter.Write(p) +} + +type debugRequestJSON struct { + Action string `json:"action"` + + KillPath string `json:"kill-path"` + KillAfter int64 `json:"kill-after"` +} + +type debugResultJSON struct { + KillAfter map[string]int64 `json:"kill-after"` +} + +func (s *Store) debugEndpoint(w http.ResponseWriter, req *http.Request) { + if req.Method == http.MethodGet { + res := debugResultJSON{ + KillAfter: s.killAfter, + } + out, err := json.MarshalIndent(res, "", " ") + if err != nil { + http.Error(w, fmt.Sprintf("cannot marshal: %v: %v", res, err), 500) + return + } + w.Write(out) + return + } + + if req.Method != http.MethodPost { + w.WriteHeader(http.StatusMethodNotAllowed) + return + } + + var debugReq *debugRequestJSON + decoder := json.NewDecoder(req.Body) + if err := decoder.Decode(&debugReq); err != nil { + http.Error(w, fmt.Sprintf("cannot decode request body: %v", err), 400) + return + } + + switch debugReq.Action { + case "kill-request": + s.debugActionKillDownload(debugReq) + default: + w.WriteHeader(400) + fmt.Fprintf(w, "unexpected debug action %q", debugReq.Action) + } +} + +func (s *Store) debugActionKillDownload(debugReq *debugRequestJSON) { + if debugReq.KillAfter == 0 { + delete(s.killAfter, debugReq.KillPath) + return + } + s.killAfter[debugReq.KillPath] = debugReq.KillAfter +} + +func logRangeHeader(handler http.HandlerFunc) http.HandlerFunc { + return func(w http.ResponseWriter, req *http.Request) { + path := req.URL.Path + if len(req.Header["Range"]) > 0 { + logger.Noticef(`requested range for %s is %v`, path, req.Header["Range"]) + } + handler(w, req) + } +} + +func (s *Store) applyKillAfter(handler http.HandlerFunc) http.HandlerFunc { + return func(w http.ResponseWriter, req *http.Request) { + path := req.URL.Path + killAfter, exists := s.killAfter[path] + if !exists { + handler(w, req) + return + } + + kaw := &killAfterWriter{ + ResponseWriter: w, + killAfter: killAfter, + } + handler(kaw, req) + + if kaw.killAfter < 0 { + logger.Noticef("%s was force killed, quota exceeded", path) + } + + // update killAfter for path after write finishes + s.killAfter[path] = kaw.killAfter + } +} + func (s *Store) searchEndpoint(w http.ResponseWriter, req *http.Request) { w.WriteHeader(501) fmt.Fprintf(w, "search not implemented") diff --git a/tests/lib/fakestore/store/store_test.go b/tests/lib/fakestore/store/store_test.go index 3534ba7bdda..1170b2e7d9d 100644 --- a/tests/lib/fakestore/store/store_test.go +++ b/tests/lib/fakestore/store/store_test.go @@ -1303,3 +1303,93 @@ func (s *storeTestSuite) TestSnapActionEndpointUnknownSnapAutoRefresh(c *C) { }, }) } + +func (s *storeTestSuite) TestDebugEndpointKillAfter(c *C) { + snapFn := s.makeTestSnap(c, "name: foo\nversion: 1") + snapInfo, err := os.Stat(snapFn) + c.Assert(err, IsNil) + + downloadPath := "/download/foo_1_all.snap" + killAfter := int64(512) + c.Assert(snapInfo.Size() > killAfter, Equals, true, + Commentf("test snap must be larger than kill-after threshold")) + + // Set a rule + resp, err := s.StorePostJSON("/debug", []byte(fmt.Sprintf(`{ + "action": "kill-request", + "kill-path": "%s", + "kill-after": %d + }`, downloadPath, killAfter))) + c.Assert(err, IsNil) + resp.Body.Close() + + resp, err = s.StoreGet("/debug") + c.Assert(err, IsNil) + defer resp.Body.Close() + + c.Assert(resp.StatusCode, Equals, 200) + var body debugResultJSON + c.Assert(json.NewDecoder(resp.Body).Decode(&body), IsNil) + c.Check(body.KillAfter, DeepEquals, map[string]int64{ + downloadPath: killAfter, + }) + + // Download is interrupted, we get fewer bytes than the full snap + resp, err = s.StoreGet(downloadPath) + c.Assert(err, IsNil) + defer resp.Body.Close() + + got, _ := io.ReadAll(resp.Body) + // Connection forcefully closed mid-transfer + c.Check(int64(len(got)) < snapInfo.Size(), Equals, true) + + // Clear it by setting kill-after to 0 + resp, err = s.StorePostJSON("/debug", []byte(fmt.Sprintf(`{ + "action": "kill-request", + "kill-path": "%s", + "kill-after": 0 + }`, downloadPath))) + c.Assert(err, IsNil) + resp.Body.Close() + + resp, err = s.StoreGet("/debug") + c.Assert(err, IsNil) + defer resp.Body.Close() + + var bodyAfterClear debugResultJSON + c.Assert(json.NewDecoder(resp.Body).Decode(&bodyAfterClear), IsNil) + c.Check(bodyAfterClear.KillAfter, HasLen, 0) + + // Download succeeds after clearing kill-after + resp, err = s.StoreGet(downloadPath) + c.Assert(err, IsNil) + defer resp.Body.Close() + + c.Assert(resp.StatusCode, Equals, 200) + got, err = io.ReadAll(resp.Body) + c.Assert(err, IsNil) + c.Check(int64(len(got)), Equals, snapInfo.Size()) +} + +func (s *storeTestSuite) TestDebugEndpointUnknownAction(c *C) { + resp, err := s.StorePostJSON("/debug", []byte(`{ + "action": "unknown-action" + }`)) + c.Assert(err, IsNil) + defer resp.Body.Close() + + c.Assert(resp.StatusCode, Equals, 400) + body, err := io.ReadAll(resp.Body) + c.Assert(err, IsNil) + c.Check(string(body), Equals, `unexpected debug action "unknown-action"`) +} + +func (s *storeTestSuite) TestDebugEndpointMethodNotAllowed(c *C) { + req, err := http.NewRequest(http.MethodPut, s.store.URL()+"/debug", nil) + c.Assert(err, IsNil) + resp, err := s.client.Do(req) + c.Assert(err, IsNil) + defer resp.Body.Close() + + c.Assert(resp.StatusCode, Equals, http.StatusMethodNotAllowed) +} From 7db32812ac1ca01eed7508819f2cf01d083d7d48 Mon Sep 17 00:00:00 2001 From: Maciej Borzecki Date: Thu, 9 Apr 2026 07:39:34 +0200 Subject: [PATCH 2/7] fixup! tests: add /debug endpoint to fakestore to allow interrupting downloads --- tests/lib/fakestore/store/store.go | 2 +- tests/lib/fakestore/store/store_test.go | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/lib/fakestore/store/store.go b/tests/lib/fakestore/store/store.go index 08c894feb19..f6f3c9a84a7 100644 --- a/tests/lib/fakestore/store/store.go +++ b/tests/lib/fakestore/store/store.go @@ -410,7 +410,7 @@ func (s *Store) debugEndpoint(w http.ResponseWriter, req *http.Request) { } if req.Method != http.MethodPost { - w.WriteHeader(http.StatusMethodNotAllowed) + w.WriteHeader(405) // Method Not Allowed return } diff --git a/tests/lib/fakestore/store/store_test.go b/tests/lib/fakestore/store/store_test.go index 1170b2e7d9d..d7c19576f3c 100644 --- a/tests/lib/fakestore/store/store_test.go +++ b/tests/lib/fakestore/store/store_test.go @@ -1391,5 +1391,5 @@ func (s *storeTestSuite) TestDebugEndpointMethodNotAllowed(c *C) { c.Assert(err, IsNil) defer resp.Body.Close() - c.Assert(resp.StatusCode, Equals, http.StatusMethodNotAllowed) + c.Assert(resp.StatusCode, Equals, 405) } From 6c3a37df9a06021108d4848bf4422b785ebc1f8a Mon Sep 17 00:00:00 2001 From: Maciej Borzecki Date: Thu, 9 Apr 2026 09:36:06 +0200 Subject: [PATCH 3/7] tests/lib/fakestore/store: add synchronization, add debug reset action Signed-off-by: Maciej Borzecki --- tests/lib/fakestore/store/store.go | 70 ++++++++++++++++++++----- tests/lib/fakestore/store/store_test.go | 39 ++++++++++++++ 2 files changed, 97 insertions(+), 12 deletions(-) diff --git a/tests/lib/fakestore/store/store.go b/tests/lib/fakestore/store/store.go index f6f3c9a84a7..e939cbabd1b 100644 --- a/tests/lib/fakestore/store/store.go +++ b/tests/lib/fakestore/store/store.go @@ -35,6 +35,7 @@ import ( "regexp" "strconv" "strings" + "sync" "time" "github.com/snapcore/snapd/asserts" @@ -67,6 +68,8 @@ type snapCachedInfo struct { // Store is our snappy software store implementation type Store struct { + lock sync.Mutex + url string blobDir string assertDir string @@ -80,6 +83,9 @@ type Store struct { snapsCache map[string]snapCachedInfo + // endpoint -> quota value, note this is stateful, i.e. the quota is counted + // for all requests to a given endpoint and after exceeding it, all + // subsequent requests will fail until it is reset though a request killAfter map[string]int64 } @@ -397,12 +403,16 @@ type debugResultJSON struct { func (s *Store) debugEndpoint(w http.ResponseWriter, req *http.Request) { if req.Method == http.MethodGet { - res := debugResultJSON{ - KillAfter: s.killAfter, - } - out, err := json.MarshalIndent(res, "", " ") + out, err := func() ([]byte, error) { + s.lock.Lock() + defer s.lock.Unlock() + res := debugResultJSON{ + KillAfter: s.killAfter, + } + return json.Marshal(res) + }() if err != nil { - http.Error(w, fmt.Sprintf("cannot marshal: %v: %v", res, err), 500) + http.Error(w, fmt.Sprintf("cannot marshal: %v", err), 500) return } w.Write(out) @@ -421,21 +431,46 @@ func (s *Store) debugEndpoint(w http.ResponseWriter, req *http.Request) { return } + var err error switch debugReq.Action { case "kill-request": - s.debugActionKillDownload(debugReq) + err = s.debugActionKillDownload(debugReq) + case "reset": + s.debugActionReset(debugReq) default: + err = fmt.Errorf("unexpected debug action %q", debugReq.Action) + } + if err != nil { w.WriteHeader(400) - fmt.Fprintf(w, "unexpected debug action %q", debugReq.Action) + fmt.Fprint(w, err.Error()) } } -func (s *Store) debugActionKillDownload(debugReq *debugRequestJSON) { +func (s *Store) debugActionKillDownload(debugReq *debugRequestJSON) error { + if debugReq.KillPath == "" { + return fmt.Errorf("kill-path cannot be empty") + } + + if strings.HasPrefix(debugReq.KillPath, "/debug/") { + return fmt.Errorf("kill-path cannot be applied to /debug/ endpoints") + } + + s.lock.Lock() + defer s.lock.Unlock() + if debugReq.KillAfter == 0 { delete(s.killAfter, debugReq.KillPath) - return + } else { + s.killAfter[debugReq.KillPath] = debugReq.KillAfter } - s.killAfter[debugReq.KillPath] = debugReq.KillAfter + return nil +} + +func (s *Store) debugActionReset(debugReq *debugRequestJSON) { + s.lock.Lock() + defer s.lock.Unlock() + + s.killAfter = map[string]int64{} } func logRangeHeader(handler http.HandlerFunc) http.HandlerFunc { @@ -451,7 +486,14 @@ func logRangeHeader(handler http.HandlerFunc) http.HandlerFunc { func (s *Store) applyKillAfter(handler http.HandlerFunc) http.HandlerFunc { return func(w http.ResponseWriter, req *http.Request) { path := req.URL.Path - killAfter, exists := s.killAfter[path] + + killAfter, exists := func() (int64, bool) { + s.lock.Lock() + defer s.lock.Unlock() + v, ok := s.killAfter[path] + return v, ok + }() + if !exists { handler(w, req) return @@ -468,7 +510,11 @@ func (s *Store) applyKillAfter(handler http.HandlerFunc) http.HandlerFunc { } // update killAfter for path after write finishes - s.killAfter[path] = kaw.killAfter + func() { + s.lock.Lock() + defer s.lock.Unlock() + s.killAfter[path] = kaw.killAfter + }() } } diff --git a/tests/lib/fakestore/store/store_test.go b/tests/lib/fakestore/store/store_test.go index d7c19576f3c..bbd9aa86151 100644 --- a/tests/lib/fakestore/store/store_test.go +++ b/tests/lib/fakestore/store/store_test.go @@ -1393,3 +1393,42 @@ func (s *storeTestSuite) TestDebugEndpointMethodNotAllowed(c *C) { c.Assert(resp.StatusCode, Equals, 405) } + +func (s *storeTestSuite) TestDebugActionReset(c *C) { + // Set a rule for endpoint connection interrupt + resp, err := s.StorePostJSON("/debug", []byte(`{ + "action": "kill-request", + "kill-path": "/foo/bar", + "kill-after": 123 + }`)) + c.Assert(err, IsNil) + resp.Body.Close() + c.Assert(resp.StatusCode, Equals, 200) + + resp, err = s.StoreGet("/debug") + c.Assert(err, IsNil) + defer resp.Body.Close() + + var buf bytes.Buffer + c.Assert(resp.StatusCode, Equals, 200) + _, err = io.Copy(&buf, resp.Body) + c.Assert(err, IsNil) + c.Check(buf.String(), Equals, `{"kill-after":{"/foo/bar":123}}`) + + // Clear it by setting kill-after to 0 + resp, err = s.StorePostJSON("/debug", []byte(`{ + "action": "reset" + }`)) + c.Assert(err, IsNil) + resp.Body.Close() + c.Assert(resp.StatusCode, Equals, 200) + + resp, err = s.StoreGet("/debug") + c.Assert(err, IsNil) + defer resp.Body.Close() + + buf.Reset() + _, err = io.Copy(&buf, resp.Body) + c.Assert(err, IsNil) + c.Check(buf.String(), Equals, `{"kill-after":{}}`) +} From 79a2a290315d1e18ee3d47a1e1d42194bdee9908 Mon Sep 17 00:00:00 2001 From: Maciej Borzecki Date: Fri, 10 Apr 2026 09:04:40 +0200 Subject: [PATCH 4/7] fixup! tests: add /debug endpoint to fakestore to allow interrupting downloads --- tests/lib/fakestore/store/store.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/lib/fakestore/store/store.go b/tests/lib/fakestore/store/store.go index e939cbabd1b..20205a31c1d 100644 --- a/tests/lib/fakestore/store/store.go +++ b/tests/lib/fakestore/store/store.go @@ -85,7 +85,7 @@ type Store struct { // endpoint -> quota value, note this is stateful, i.e. the quota is counted // for all requests to a given endpoint and after exceeding it, all - // subsequent requests will fail until it is reset though a request + // subsequent requests will fail until it is reset through a request killAfter map[string]int64 } From c55bc94a5baa46b1dc5ced90c6938442fc38efce Mon Sep 17 00:00:00 2001 From: Maciej Borzecki Date: Fri, 10 Apr 2026 09:04:56 +0200 Subject: [PATCH 5/7] tests/lib/fakestore/store: close the connection after exceeding the limit Signed-off-by: Maciej Borzecki --- tests/lib/fakestore/store/store.go | 43 ++++++++++++++++++------- tests/lib/fakestore/store/store_test.go | 4 +-- 2 files changed, 34 insertions(+), 13 deletions(-) diff --git a/tests/lib/fakestore/store/store.go b/tests/lib/fakestore/store/store.go index 20205a31c1d..7765053a927 100644 --- a/tests/lib/fakestore/store/store.go +++ b/tests/lib/fakestore/store/store.go @@ -373,21 +373,42 @@ type killAfterWriter struct { } func (kaw *killAfterWriter) Write(p []byte) (int, error) { - if kaw.killAfter >= 0 { - kaw.killAfter -= int64(len(p)) - } - if kaw.killAfter < 0 { - // hijack the connection to force a hard drop - hj, ok := kaw.ResponseWriter.(http.Hijacker) - if ok { - conn, _, _ := hj.Hijack() - conn.Close() // hard close the TCP connection - } + // already exceeded the quota, kill immediately + kaw.hijackAndClose() return 0, fmt.Errorf("connection killed") } - return kaw.ResponseWriter.Write(p) + toWrite := p + shouldKill := false + if int64(len(p)) > kaw.killAfter { + // write only up to the remaining quota + toWrite = p[:kaw.killAfter] + shouldKill = true + } + + n, err := kaw.ResponseWriter.Write(toWrite) + kaw.killAfter -= int64(n) + + if shouldKill { + kaw.hijackAndClose() + return n, fmt.Errorf("connection killed") + } + + return n, err +} + +func (kaw *killAfterWriter) hijackAndClose() { + // flush any buffered data before closing + if f, ok := kaw.ResponseWriter.(http.Flusher); ok { + f.Flush() + } + // and proceed to close + hj, ok := kaw.ResponseWriter.(http.Hijacker) + if ok { + conn, _, _ := hj.Hijack() + conn.Close() + } } type debugRequestJSON struct { diff --git a/tests/lib/fakestore/store/store_test.go b/tests/lib/fakestore/store/store_test.go index bbd9aa86151..8517546a68d 100644 --- a/tests/lib/fakestore/store/store_test.go +++ b/tests/lib/fakestore/store/store_test.go @@ -1340,8 +1340,8 @@ func (s *storeTestSuite) TestDebugEndpointKillAfter(c *C) { defer resp.Body.Close() got, _ := io.ReadAll(resp.Body) - // Connection forcefully closed mid-transfer - c.Check(int64(len(got)) < snapInfo.Size(), Equals, true) + // Connection forcefully closed mid-transfer, exactly killAfter bytes received + c.Check(int64(len(got)), Equals, killAfter) // Clear it by setting kill-after to 0 resp, err = s.StorePostJSON("/debug", []byte(fmt.Sprintf(`{ From 4e3a53130f23373bdaf55541fdf46d5f3655b7b8 Mon Sep 17 00:00:00 2001 From: Maciej Borzecki Date: Fri, 10 Apr 2026 12:32:17 +0200 Subject: [PATCH 6/7] fixup! tests/lib/fakestore/store: close the connection after exceeding the limit --- tests/lib/fakestore/store/store.go | 2 +- tests/lib/fakestore/store/store_test.go | 10 ++++++++++ 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/tests/lib/fakestore/store/store.go b/tests/lib/fakestore/store/store.go index 7765053a927..260aec4476f 100644 --- a/tests/lib/fakestore/store/store.go +++ b/tests/lib/fakestore/store/store.go @@ -373,7 +373,7 @@ type killAfterWriter struct { } func (kaw *killAfterWriter) Write(p []byte) (int, error) { - if kaw.killAfter < 0 { + if kaw.killAfter <= 0 { // already exceeded the quota, kill immediately kaw.hijackAndClose() return 0, fmt.Errorf("connection killed") diff --git a/tests/lib/fakestore/store/store_test.go b/tests/lib/fakestore/store/store_test.go index 8517546a68d..05e082ca074 100644 --- a/tests/lib/fakestore/store/store_test.go +++ b/tests/lib/fakestore/store/store_test.go @@ -1343,6 +1343,16 @@ func (s *storeTestSuite) TestDebugEndpointKillAfter(c *C) { // Connection forcefully closed mid-transfer, exactly killAfter bytes received c.Check(int64(len(got)), Equals, killAfter) + // Retry the request, which should be killed after receiving 0 bytes because + // the killAfter effect is stateful. + resp, err = s.StoreGet(downloadPath) + c.Assert(err, IsNil) + defer resp.Body.Close() + + got, _ = io.ReadAll(resp.Body) + // Connection forcefully closed mid-transfer, exactly killAfter bytes received + c.Check(int64(len(got)), Equals, int64(0)) + // Clear it by setting kill-after to 0 resp, err = s.StorePostJSON("/debug", []byte(fmt.Sprintf(`{ "action": "kill-request", From ee7319d7a96108506f8aecde8f6e3746608c2edb Mon Sep 17 00:00:00 2001 From: Maciej Borzecki Date: Fri, 10 Apr 2026 17:49:53 +0200 Subject: [PATCH 7/7] tests/lib/fakestore/store: fix race in how the quota is counted Fix a race in consumign and trackign the left quota. Signed-off-by: Maciej Borzecki --- tests/lib/fakestore/store/store.go | 59 ++++++++++++++++++------------ 1 file changed, 35 insertions(+), 24 deletions(-) diff --git a/tests/lib/fakestore/store/store.go b/tests/lib/fakestore/store/store.go index 260aec4476f..03d13111004 100644 --- a/tests/lib/fakestore/store/store.go +++ b/tests/lib/fakestore/store/store.go @@ -369,28 +369,25 @@ type detailsReplyJSON struct { type killAfterWriter struct { http.ResponseWriter - killAfter int64 + path string + consumeQuota func(want int) int } func (kaw *killAfterWriter) Write(p []byte) (int, error) { - if kaw.killAfter <= 0 { - // already exceeded the quota, kill immediately - kaw.hijackAndClose() - return 0, fmt.Errorf("connection killed") - } - toWrite := p shouldKill := false - if int64(len(p)) > kaw.killAfter { + + got := kaw.consumeQuota(len(toWrite)) + if len(p) > got { // write only up to the remaining quota - toWrite = p[:kaw.killAfter] + toWrite = p[:got] shouldKill = true } n, err := kaw.ResponseWriter.Write(toWrite) - kaw.killAfter -= int64(n) if shouldKill { + logger.Noticef("request to %s was force killed, quota exceeded", kaw.path) kaw.hijackAndClose() return n, fmt.Errorf("connection killed") } @@ -508,11 +505,11 @@ func (s *Store) applyKillAfter(handler http.HandlerFunc) http.HandlerFunc { return func(w http.ResponseWriter, req *http.Request) { path := req.URL.Path - killAfter, exists := func() (int64, bool) { + exists := func() bool { s.lock.Lock() defer s.lock.Unlock() - v, ok := s.killAfter[path] - return v, ok + _, ok := s.killAfter[path] + return ok }() if !exists { @@ -522,20 +519,34 @@ func (s *Store) applyKillAfter(handler http.HandlerFunc) http.HandlerFunc { kaw := &killAfterWriter{ ResponseWriter: w, - killAfter: killAfter, + path: path, + consumeQuota: func(want int) int { + s.lock.Lock() + defer s.lock.Unlock() + + v, ok := s.killAfter[path] + if !ok { + // no quota set + return want + } + + left := int(v) + + var got int + if want > left { + got = left + left = 0 + } else { + got = want + left -= want + } + s.killAfter[path] = int64(left) + + return got + }, } handler(kaw, req) - if kaw.killAfter < 0 { - logger.Noticef("%s was force killed, quota exceeded", path) - } - - // update killAfter for path after write finishes - func() { - s.lock.Lock() - defer s.lock.Unlock() - s.killAfter[path] = kaw.killAfter - }() } }