fix: flush downstream ResponseWriter on upstream watch stream close to ensure clean EOF delivery

pkg/yurthub/cachemanager/cache_manager.go

@@ -72,6 +72,11 @@ type CacheResult struct {
 	Msg    string
 }
 
+type pendingWrite struct {
+	key storage.Key
+	obj runtime.Object
+}
+
 type cacheManager struct {
 	sync.RWMutex
 	storage               StorageWrapper
@@ -80,6 +85,7 @@ type cacheManager struct {
 	configManager         *configuration.Manager
 	listSelectorCollector map[storage.Key]string
 	inMemoryCache         map[string]runtime.Object
+	pendingWrites         []pendingWrite
 }
 
 // NewCacheManager creates a new CacheManager
@@ -96,6 +102,7 @@ func NewCacheManager(
 		configManager:         configManager,
 		listSelectorCollector: make(map[storage.Key]string),
 		inMemoryCache:         make(map[string]runtime.Object),
+		pendingWrites:         make([]pendingWrite, 0),
 	}
 	return cm
 }
@@ -556,11 +563,35 @@ func (cm *cacheManager) saveListObject(ctx context.Context, info *apirequest.Req
 			objs[key] = items[i]
 		}
 		// if no objects in cloud cluster(objs is empty), it will clean the old files in the path of rootkey
-		return cm.storage.ReplaceComponentList(comp, schema.GroupVersionResource{
+		err := cm.storage.ReplaceComponentList(comp, schema.GroupVersionResource{
 			Group:    info.APIGroup,
 			Version:  info.APIVersion,
 			Resource: info.Resource,
 		}, info.Namespace, objs)
+		if err != nil {
+			return err
+		}
+		cm.drainPendingWrites()
+		return nil
+	}
+}
+
+func (cm *cacheManager) enqueuePendingWrite(key storage.Key, obj runtime.Object) {
+	cm.Lock()
+	defer cm.Unlock()
+	cm.pendingWrites = append(cm.pendingWrites, pendingWrite{key: key, obj: obj})
+}
+
+func (cm *cacheManager) drainPendingWrites() {
+	cm.Lock()
+	writes := cm.pendingWrites
+	cm.pendingWrites = nil
+	cm.Unlock()
+
+	for _, pw := range writes {
+		if err := cm.storeObjectWithKey(pw.key, pw.obj); err != nil {
+			klog.Errorf("could not store pending object %s, %v", pw.key.Key(), err)
+		}
 	}
 }
 
@@ -679,12 +710,14 @@ func (cm *cacheManager) storeObjectWithKey(key storage.Key, obj runtime.Object)
 		if err := cm.storage.Create(key, obj); err != nil {
 			if errors.Is(err, storage.ErrStorageAccessConflict) {
 				klog.V(2).Infof("skip to cache obj because key(%s) is under processing", key.Key())
+				cm.enqueuePendingWrite(key, obj)
 				return nil
 			}
 			return fmt.Errorf("could not create obj of key: %s, %v", key.Key(), err)
 		}
 	case errors.Is(err, storage.ErrStorageAccessConflict):
 		klog.V(2).Infof("skip to cache watch event because key(%s) is under processing", key.Key())
+		cm.enqueuePendingWrite(key, obj)
 		return nil
 	default:
 		return fmt.Errorf("could not store obj with rv %s of key: %s, %v", newRv, key.Key(), err)

pkg/yurthub/cachemanager/cache_manager_test.go

@@ -19,13 +19,15 @@ package cachemanager
 import (
 	"bytes"
 	"context"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"io"
 	"net/http"
 	"net/http/httptest"
 	"os"
 	"strings"
+	"sync"
 	"testing"
 	"time"
 
@@ -3454,4 +3456,189 @@ func TestIsListRequestWithNameFieldSelector(t *testing.T) {
 	}
 }
 
-// TODO: in-memory cache unit tests
+type blockingStorageWrapper struct {
+	StorageWrapper
+	replaceStartedCh chan struct{}
+	replaceUnblockCh chan struct{}
+	replaceConflict  sync.Mutex
+	inConflict       bool
+}
+
+func (b *blockingStorageWrapper) ReplaceComponentList(component string, gvr schema.GroupVersionResource, namespace string, objs map[storage.Key]runtime.Object) error {
+	b.replaceStartedCh <- struct{}{}
+	<-b.replaceUnblockCh
+	return b.StorageWrapper.ReplaceComponentList(component, gvr, namespace, objs)
+}
+
+func (b *blockingStorageWrapper) Update(key storage.Key, obj runtime.Object, rv uint64) (runtime.Object, error) {
+	b.replaceConflict.Lock()
+	conflict := b.inConflict
+	b.replaceConflict.Unlock()
+	if conflict {
+		return nil, storage.ErrStorageAccessConflict
+	}
+	return b.StorageWrapper.Update(key, obj, rv)
+}
+
+func (b *blockingStorageWrapper) Create(key storage.Key, obj runtime.Object) error {
+	b.replaceConflict.Lock()
+	conflict := b.inConflict
+	b.replaceConflict.Unlock()
+	if conflict {
+		return storage.ErrStorageAccessConflict
+	}
+	return b.StorageWrapper.Create(key, obj)
+}
+
+func TestWatchDuringReplaceComponentListDrainsPendingWritesToDisk(t *testing.T) {
+	tempDir := t.TempDir()
+	dStorage, err := disk.NewDiskStorage(tempDir)
+	if err != nil {
+		t.Fatalf("could not create disk storage, %v", err)
+	}
+	sWrapper := NewStorageWrapper(dStorage)
+
+	bWrapper := &blockingStorageWrapper{
+		StorageWrapper:   sWrapper,
+		replaceStartedCh: make(chan struct{}),
+		replaceUnblockCh: make(chan struct{}),
+	}
+
+	fakeSharedInformerFactory := informers.NewSharedInformerFactory(fake.NewSimpleClientset(), 0)
+	configManager := configuration.NewConfigurationManager("node1", fakeSharedInformerFactory)
+	serializerM := serializer.NewSerializerManager()
+	restMapperM, err := hubmeta.NewRESTMapperManager(tempDir)
+	if err != nil {
+		t.Fatalf("could not create RESTMapperManager, %v", err)
+	}
+
+	cm := NewCacheManager(bWrapper, serializerM, restMapperM, configManager)
+
+	pod1 := &v1.Pod{
+		TypeMeta: metav1.TypeMeta{
+			APIVersion: "v1",
+			Kind:       "Pod",
+		},
+		ObjectMeta: metav1.ObjectMeta{
+			Name:            "pod1",
+			Namespace:       "default",
+			ResourceVersion: "1",
+		},
+		Spec: v1.PodSpec{
+			NodeName: "node1",
+		},
+	}
+	podList := &v1.PodList{
+		TypeMeta: metav1.TypeMeta{
+			APIVersion: "v1",
+			Kind:       "PodList",
+		},
+		ListMeta: metav1.ListMeta{
+			ResourceVersion: "1",
+		},
+		Items: []v1.Pod{*pod1},
+	}
+
+	podListBytes, err := json.Marshal(podList)
+	if err != nil {
+		t.Fatalf("failed to marshal pod list, %v", err)
+	}
+
+	pod2 := &v1.Pod{
+		TypeMeta: metav1.TypeMeta{
+			APIVersion: "v1",
+			Kind:       "Pod",
+		},
+		ObjectMeta: metav1.ObjectMeta{
+			Name:            "pod2",
+			Namespace:       "default",
+			ResourceVersion: "2",
+		},
+		Spec: v1.PodSpec{
+			NodeName: "node1",
+		},
+	}
+
+	resolver := newTestRequestInfoResolver()
+
+	var listErr error
+	var listWg sync.WaitGroup
+	listWg.Add(1)
+	go func() {
+		defer listWg.Done()
+		req, _ := http.NewRequest("GET", "/api/v1/namespaces/default/pods", nil)
+		req.Header.Set("User-Agent", "kubelet")
+		req.Header.Set("Accept", "application/json")
+		req.RemoteAddr = "127.0.0.1"
+
+		var handler http.Handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			ctx := util.WithRespContentType(r.Context(), "application/json")
+			r = r.WithContext(ctx)
+			listErr = cm.CacheResponse(r, io.NopCloser(bytes.NewReader(podListBytes)), nil)
+		})
+		handler = proxyutil.WithListRequestSelector(handler)
+		handler = proxyutil.WithRequestClientComponent(handler)
+		handler = filters.WithRequestInfo(handler, resolver)
+		handler.ServeHTTP(httptest.NewRecorder(), req)
+	}()
+
+	<-bWrapper.replaceStartedCh
+	bWrapper.replaceConflict.Lock()
+	bWrapper.inConflict = true
+	bWrapper.replaceConflict.Unlock()
+
+	s := serializerM.CreateSerializer("application/json", "", "v1", "pods")
+	eventBuf := new(bytes.Buffer)
+	if _, err := s.WatchEncode(eventBuf, &watch.Event{
+		Type:   watch.Added,
+		Object: pod2,
+	}); err != nil {
+		t.Fatalf("failed to encode watch event, %v", err)
+	}
+
+	watchReq, _ := http.NewRequest("GET", "/api/v1/namespaces/default/pods?watch=true", nil)
+	watchReq.Header.Set("User-Agent", "kubelet")
+	watchReq.Header.Set("Accept", "application/json")
+	watchReq.RemoteAddr = "127.0.0.1"
+
+	var watchErr error
+	var watchHandler http.Handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		ctx := util.WithRespContentType(r.Context(), "application/json")
+		r = r.WithContext(ctx)
+		watchErr = cm.CacheResponse(r, io.NopCloser(bytes.NewReader(eventBuf.Bytes())), nil)
+	})
+	watchHandler = proxyutil.WithRequestClientComponent(watchHandler)
+	watchHandler = filters.WithRequestInfo(watchHandler, resolver)
+	watchHandler.ServeHTTP(httptest.NewRecorder(), watchReq)
+
+	bWrapper.replaceConflict.Lock()
+	bWrapper.inConflict = false
+	bWrapper.replaceConflict.Unlock()
+	bWrapper.replaceUnblockCh <- struct{}{}
+
+	listWg.Wait()
+
+	if listErr != nil {
+		t.Errorf("unexpected list error: %v", listErr)
+	}
+	if watchErr != nil && !errors.Is(watchErr, io.EOF) {
+		t.Errorf("unexpected watch error: %v", watchErr)
+	}
+
+	pod2Key, err := sWrapper.KeyFunc(storage.KeyBuildInfo{
+		Component: "kubelet",
+		Namespace: "default",
+		Name:      "pod2",
+		Resources: "pods",
+		Group:     "",
+		Version:   "v1",
+	})
+	if err != nil {
+		t.Fatalf("failed to build key for pod2, %v", err)
+	}
+
+	diskObj, err := sWrapper.Get(pod2Key)
+	if err != nil || diskObj == nil {
+		t.Fatalf("expected pod2 to be present on disk, but got err: %v, obj: %v", err, diskObj)
+	}
+}

pkg/yurthub/proxy/remote/remote.go

@@ -122,6 +122,9 @@ func (rp *RemoteProxy) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
 	}
 
 	rp.reverseProxy.ServeHTTP(rw, req)
+	if flusher, ok := rw.(http.Flusher); ok {
+		flusher.Flush()
+	}
 }
 
 // RoundTrip is used to implement http.RoundTripper for RemoteProxy.