export: make the cancelled-resume test deterministic

internal/export/export_test.go

@@ -104,6 +104,29 @@ func (f *fixture) downloads() []testserver.Event {
 	return out
 }
 
+// partSizes maps each archive of example.com that has a partial file
+// ("images-1.zip") to that file's size. It reads the size through an open
+// file, which is current on every platform even while a download still
+// writes to it.
+func (f *fixture) partSizes() map[string]int64 {
+	out := map[string]int64{}
+	paths, _ := filepath.Glob(filepath.Join(f.dir, "example.com", "*"+partSuffix))
+	for _, p := range paths {
+		name := filepath.Base(p)
+		i := strings.Index(name, ".zip.")
+		file, err := os.Open(p)
+		if i < 0 || err != nil {
+			continue
+		}
+		st, err := file.Stat()
+		file.Close()
+		if err == nil {
+			out[name[:i+len(".zip")]] = st.Size()
+		}
+	}
+	return out
+}
+
 func TestRunDownloadsArchivesWhileTheExportRuns(t *testing.T) {
 	f := newFixture(t)
 	a1, a2, d1 := zipOf(t, "example.com/a.jpg"), zipOf(t, "example.com/b.jpg"), zipOf(t, "c.png")
@@ -759,28 +782,39 @@ func TestRunResumesDownloadsCancelledMidway(t *testing.T) {
 	a, b := zipOf(t, "a"), zipOf(t, "b")
 	id := f.srv.AddExport(api.ExportCompleted, testserver.ExportDomain{Domain: "example.com", Archives: [][]byte{a, b}})
 	f.srv.Inject(testserver.RouteDownload, testserver.Response{StallAfter: 30}, testserver.Response{StallAfter: 40})
+	// Whichever archive starts first stalls after 30 bytes, the other after
+	// 40. The run is cancelled once both partial files hold those bytes,
+	// however long the downloads take to get there.
+	stalled := func(parts map[string]int64) bool {
+		one, two := parts["images-1.zip"], parts["images-2.zip"]
+		return len(parts) == 2 && (one == 30 && two == 40 || one == 40 && two == 30)
+	}
 	ctx, cancel := context.WithCancel(context.Background())
-	stalled := 0
-	f.srv.Hook = func(route testserver.Route, _ *http.Request) {
-		if route == testserver.RouteDownload {
-			f.mu.Lock()
-			stalled++
-			if stalled == 2 {
-				time.AfterFunc(100*time.Millisecond, cancel)
+	defer cancel()
+	ran := make(chan struct{})
+	go func() {
+		defer cancel()
+		for deadline := time.Now().Add(10 * time.Second); time.Now().Before(deadline) && !stalled(f.partSizes()); {
+			select {
+			case <-ran:
+				return
+			case <-time.After(time.Millisecond):
 			}
-			f.mu.Unlock()
 		}
-	}
+	}()
 	out := f.run(ctx, id)
+	close(ran)
 	if !errors.Is(out.Err, context.Canceled) || out.Failed != 0 || out.Downloaded != 0 {
 		t.Fatalf("cancelled run: %+v", out)
 	}
-	parts, _ := filepath.Glob(filepath.Join(f.dir, "example.com", "*"+partSuffix))
-	if len(parts) != 2 {
+	parts := f.partSizes()
+	if !stalled(parts) {
 		t.Fatalf("partial files: %v", parts)
 	}
 
-	f.srv.Hook = nil
+	// The fake files the cancelled downloads only once it notices the
+	// client left, which can be after Run returned.
+	f.srv.WaitIdle(t)
 	before := len(f.downloads())
 	out = f.run(context.Background(), id)
 	if out.Err != nil || out.Downloaded != 2 {
@@ -789,9 +823,14 @@ func TestRunResumesDownloadsCancelledMidway(t *testing.T) {
 	if !bytes.Equal(f.file("example.com/images-1.zip"), a) || !bytes.Equal(f.file("example.com/images-2.zip"), b) {
 		t.Error("resumed archives differ")
 	}
-	for _, d := range f.downloads()[before:] {
-		if d.Range != "bytes=30-" && d.Range != "bytes=40-" {
-			t.Errorf("resume without the right range: %+v", d)
+	resumed := f.downloads()[before:]
+	if len(resumed) != 2 {
+		t.Fatalf("resumed downloads: %+v", resumed)
+	}
+	for _, d := range resumed {
+		archive := "images-" + d.Key[strings.LastIndex(d.Key, "/")+1:] + ".zip"
+		if want := fmt.Sprintf("bytes=%d-", parts[archive]); d.Range != want {
+			t.Errorf("%s resumed with Range %q, want %q: %+v", archive, d.Range, want, d)
 		}
 	}
 }

internal/testserver/testserver.go

@@ -171,6 +171,10 @@ type Server struct {
 	violations  []string
 	inflightPut int
 	maxPut      int
+	// active counts the requests being handled; idle is closed once it
+	// drops to zero (see WaitIdle).
+	active int
+	idle   chan struct{}
 
 	exports        map[int64]*fakeExport
 	nextExport     int64
@@ -196,14 +200,14 @@ func New(t testing.TB) *Server {
 	api.HandleFunc("/api/folders", s.handleFolders)
 	api.HandleFunc("/api/asset-exports", s.handleExports)
 	api.HandleFunc("/api/asset-exports/", s.handleExport)
-	s.API = httptest.NewServer(api)
-	s.S3 = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+	s.API = httptest.NewServer(s.track(api))
+	s.S3 = httptest.NewServer(s.track(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		if strings.HasPrefix(r.URL.Path, "/exports/") {
 			s.handleDownload(w, r)
 			return
 		}
 		s.handlePut(w, r)
-	}))
+	})))
 	s.URL = s.API.URL + "/api"
 	t.Cleanup(func() {
 		s.API.Close()
@@ -270,6 +274,47 @@ func (s *Server) Events() []Event {
 	return append([]Event(nil), s.events...)
 }
 
+// WaitIdle waits until no request is being handled, so Events holds every
+// request whose handler has started. A request is filed when its handler
+// ends, and the handler of a request the client gave up on (a stalled
+// download, say) ends only once the fake notices, which can be after the
+// client has moved on. The test fails if requests are still being handled
+// after 10 seconds.
+func (s *Server) WaitIdle(t testing.TB) {
+	t.Helper()
+	s.mu.Lock()
+	active, idle := s.active, s.idle
+	s.mu.Unlock()
+	if active == 0 {
+		return
+	}
+	select {
+	case <-idle:
+	case <-time.After(10 * time.Second):
+		t.Fatal("testserver: requests still being handled after 10s")
+	}
+}
+
+// track counts the requests next is handling, for WaitIdle.
+func (s *Server) track(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		s.mu.Lock()
+		if s.active == 0 {
+			s.idle = make(chan struct{})
+		}
+		s.active++
+		s.mu.Unlock()
+		defer func() {
+			s.mu.Lock()
+			if s.active--; s.active == 0 {
+				close(s.idle)
+			}
+			s.mu.Unlock()
+		}()
+		next.ServeHTTP(w, r)
+	})
+}
+
 // Count returns how many requests reached a route.
 func (s *Server) Count(route Route) int {
 	n := 0