[ws-daemon] Dispose all handlers running triggered by the "Dispatch" abstraction (fueled by updates from a separate PodInformer)

Author: geropl

components/ws-daemon/pkg/cgroup/cgroup.go

@@ -79,8 +79,11 @@ func (host *PluginHost) WorkspaceAdded(ctx context.Context, ws *dispatch.Workspa
 		return xerrors.Errorf("no dispatch available")
 	}
 
-	cgroupPath, err := disp.Runtime.ContainerCGroupPath(context.Background(), ws.ContainerID)
+	cgroupPath, err := disp.Runtime.ContainerCGroupPath(ctx, ws.ContainerID)
 	if err != nil {
+		if err == context.Canceled {
+			return nil
+		}
 		return xerrors.Errorf("cannot get cgroup path for container %s: %w", ws.ContainerID, err)
 	}
 
@@ -95,8 +98,11 @@ func (host *PluginHost) WorkspaceAdded(ctx context.Context, ws *dispatch.Workspa
 		if plg.Type() != host.CGroupVersion {
 			continue
 		}
+		dispatch.GetDispatchWaitGroup(ctx).Add(1)
 
 		go func(plg Plugin) {
+			defer dispatch.GetDispatchWaitGroup(ctx).Done()
+
 			err := plg.Apply(ctx, opts)
 			if err == context.Canceled || err == context.DeadlineExceeded {
 				err = nil

components/ws-daemon/pkg/controller/workspace_operations.go

@@ -21,6 +21,7 @@ import (
 	"github.com/gitpod-io/gitpod/content-service/pkg/logs"
 	"github.com/gitpod-io/gitpod/content-service/pkg/storage"
 	"github.com/gitpod-io/gitpod/ws-daemon/pkg/content"
+	"github.com/gitpod-io/gitpod/ws-daemon/pkg/dispatch"
 	"github.com/gitpod-io/gitpod/ws-daemon/pkg/internal/session"
 	"github.com/opentracing/opentracing-go"
 	"github.com/prometheus/client_golang/prometheus"
@@ -81,6 +82,7 @@ type DefaultWorkspaceOperations struct {
 	provider               *WorkspaceProvider
 	backupWorkspaceLimiter chan struct{}
 	metrics                *Metrics
+	dispatch               *dispatch.Dispatch
 }
 
 var _ WorkspaceOperations = (*DefaultWorkspaceOperations)(nil)
@@ -106,7 +108,7 @@ type BackupOptions struct {
 	SkipBackupContent bool
 }
 
-func NewWorkspaceOperations(config content.Config, provider *WorkspaceProvider, reg prometheus.Registerer) (WorkspaceOperations, error) {
+func NewWorkspaceOperations(config content.Config, provider *WorkspaceProvider, reg prometheus.Registerer, dispatch *dispatch.Dispatch) (WorkspaceOperations, error) {
 	waitingTimeHist, waitingTimeoutCounter, err := registerConcurrentBackupMetrics(reg, "_mk2")
 	if err != nil {
 		return nil, err
@@ -121,6 +123,7 @@ func NewWorkspaceOperations(config content.Config, provider *WorkspaceProvider,
 		},
 		// we permit five concurrent backups at any given time, hence the five in the channel
 		backupWorkspaceLimiter: make(chan struct{}, 5),
+		dispatch:               dispatch,
 	}, nil
 }
 
@@ -185,6 +188,8 @@ func (wso *DefaultWorkspaceOperations) InitWorkspace(ctx context.Context, option
 		return "cannot persist workspace", err
 	}
 
+	glog.WithFields(ws.OWI()).Debug("content init done")
+
 	return "", nil
 }
 
@@ -305,6 +310,10 @@ func (wso *DefaultWorkspaceOperations) WipeWorkspace(ctx context.Context, instan
 		return err
 	}
 
+	// dispose all running "dispatch handlers", e.g. all code running on the "pod informer"-triggered part of ws-daemon
+	wso.dispatch.DisposeWorkspace(ctx, instanceID)
+
+	// remove the reference from the WorkspaceProvider, e.g. the "workspace controller" part of ws-daemon
 	wso.provider.Remove(ctx, instanceID)
 
 	return nil

components/ws-daemon/pkg/cpulimit/dispatch.go

@@ -179,8 +179,11 @@ func (d *DispatchListener) WorkspaceAdded(ctx context.Context, ws *dispatch.Work
 		return xerrors.Errorf("no dispatch available")
 	}
 
-	cgroupPath, err := disp.Runtime.ContainerCGroupPath(context.Background(), ws.ContainerID)
+	cgroupPath, err := disp.Runtime.ContainerCGroupPath(ctx, ws.ContainerID)
 	if err != nil {
+		if dispatch.IsCancelled(ctx) {
+			return nil
+		}
 		return xerrors.Errorf("cannot start governer: %w", err)
 	}
 
@@ -194,7 +197,11 @@ func (d *DispatchListener) WorkspaceAdded(ctx context.Context, ws *dispatch.Work
 		OWI:         ws.OWI(),
 		Annotations: ws.Pod.Annotations,
 	}
+
+	dispatch.GetDispatchWaitGroup(ctx).Add(1)
 	go func() {
+		defer dispatch.GetDispatchWaitGroup(ctx).Done()
+
 		<-ctx.Done()
 
 		d.mu.Lock()

components/ws-daemon/pkg/daemon/daemon.go

@@ -208,7 +208,12 @@ func NewDaemon(config Config) (*Daemon, error) {
 		config.CPULimit.CGroupBasePath,
 	)
 
-	workspaceOps, err := controller.NewWorkspaceOperations(contentCfg, controller.NewWorkspaceProvider(contentCfg.WorkingArea, hooks), wrappedReg)
+	dsptch, err := dispatch.NewDispatch(containerRuntime, clientset, config.Runtime.KubernetesNamespace, nodename, listener...)
+	if err != nil {
+		return nil, err
+	}
+
+	workspaceOps, err := controller.NewWorkspaceOperations(contentCfg, controller.NewWorkspaceProvider(contentCfg.WorkingArea, hooks), wrappedReg, dsptch)
 	if err != nil {
 		return nil, err
 	}
@@ -233,11 +238,6 @@ func NewDaemon(config Config) (*Daemon, error) {
 	housekeeping := controller.NewHousekeeping(contentCfg.WorkingArea, 5*time.Minute)
 	go housekeeping.Start(context.Background())
 
-	dsptch, err := dispatch.NewDispatch(containerRuntime, clientset, config.Runtime.KubernetesNamespace, nodename, listener...)
-	if err != nil {
-		return nil, err
-	}
-
 	dsk := diskguard.FromConfig(config.DiskSpaceGuard, clientset, nodename)
 
 	return &Daemon{

components/ws-daemon/pkg/daemon/markunmount.go

@@ -101,8 +101,24 @@ func (c *MarkUnmountFallback) WorkspaceUpdated(ctx context.Context, ws *dispatch
 	}
 	ttl := time.Duration(gracePeriod)*time.Second + propagationGracePeriod
 
+	dispatch.GetDispatchWaitGroup(ctx).Add(1)
 	go func() {
-		time.Sleep(ttl)
+		defer dispatch.GetDispatchWaitGroup(ctx).Done()
+
+		defer func() {
+			// We expect the container to be gone now. Don't keep its referenec in memory.
+			c.mu.Lock()
+			delete(c.handled, ws.InstanceID)
+			c.mu.Unlock()
+		}()
+
+		wait := time.NewTicker(ttl)
+		defer wait.Stop()
+		select {
+		case <-ctx.Done():
+			return
+		case <-wait.C:
+		}
 
 		dsp := dispatch.GetFromContext(ctx)
 		if !dsp.WorkspaceExistsOnNode(ws.InstanceID) {
@@ -111,17 +127,12 @@ func (c *MarkUnmountFallback) WorkspaceUpdated(ctx context.Context, ws *dispatch
 		}
 
 		err := unmountMark(ws.InstanceID)
-		if err != nil {
+		if err != nil && !dispatch.IsCancelled(ctx) {
 			log.WithFields(ws.OWI()).WithError(err).Error("cannot unmount mark mount from within ws-daemon")
 			c.activityCounter.WithLabelValues("false").Inc()
 		} else {
 			c.activityCounter.WithLabelValues("true").Inc()
 		}
-
-		// We expect the container to be gone now. Don't keep its referenec in memory.
-		c.mu.Lock()
-		delete(c.handled, ws.InstanceID)
-		c.mu.Unlock()
 	}()
 
 	return nil

components/ws-daemon/pkg/dispatch/dispatch.go

@@ -86,6 +86,9 @@ type workspaceState struct {
 	Context        context.Context
 	Cancel         context.CancelFunc
 	Workspace      *Workspace
+
+	// this WaitGroup keeps track of when each handler is finished. It's only relied upon in DisposeWorkspace() to determine when work on a given instanceID has commenced.
+	HandlerWaitGroup sync.WaitGroup
 }
 
 type contextKey struct{}
@@ -99,6 +102,20 @@ func GetFromContext(ctx context.Context) *Dispatch {
 	return ctx.Value(contextDispatch).(*Dispatch)
 }
 
+type dispacthHandlerWaitGroupKey struct{}
+
+var (
+	contextDispatchWaitGroup = dispacthHandlerWaitGroupKey{}
+)
+
+func GetDispatchWaitGroup(ctx context.Context) *sync.WaitGroup {
+	return ctx.Value(contextDispatchWaitGroup).(*sync.WaitGroup)
+}
+
+func IsCancelled(ctx context.Context) bool {
+	return context.Cause(ctx) != nil
+}
+
 // Start starts the dispatch
 func (d *Dispatch) Start() error {
 	ifac := informers.NewSharedInformerFactoryWithOptions(d.Kubernetes, podInformerResyncInterval, informers.WithNamespace(d.KubernetesNamespace))
@@ -170,6 +187,29 @@ func (d *Dispatch) WorkspaceExistsOnNode(instanceID string) (ok bool) {
 	return
 }
 
+// DisposeWorkspace makes sure
+func (d *Dispatch) DisposeWorkspace(ctx context.Context, instanceID string) {
+	d.mu.Lock()
+	defer d.mu.Unlock()
+
+	log.WithField("instanceID", instanceID).Debugf("WS DISPOSE: %s", instanceID)
+
+	// If we have that instanceID present, cancel it's context
+	state, present := d.ctxs[instanceID]
+	if !present {
+		return
+	}
+	if state.Cancel != nil {
+		state.Cancel()
+	}
+
+	// ...and wait for all long-running/async processes/go-routines to finish
+	state.HandlerWaitGroup.Wait()
+
+	delete(d.ctxs, instanceID)
+	log.WithField("instanceID", instanceID).Debugf("WS DISPOSE DONE: %s", instanceID)
+}
+
 func (d *Dispatch) handlePodUpdate(oldPod, newPod *corev1.Pod) {
 	workspaceID, ok := newPod.Labels[wsk8s.MetaIDLabel]
 	if !ok {
@@ -182,6 +222,7 @@ func (d *Dispatch) handlePodUpdate(oldPod, newPod *corev1.Pod) {
 	if d.NodeName != "" && newPod.Spec.NodeName != d.NodeName {
 		return
 	}
+	log.WithField("instanceID", workspaceInstanceID).Debugf("POD UPDATE: %s", workspaceInstanceID)
 
 	d.mu.Lock()
 	defer d.mu.Unlock()
@@ -190,19 +231,21 @@ func (d *Dispatch) handlePodUpdate(oldPod, newPod *corev1.Pod) {
 	if !ok {
 		// we haven't seen this pod before - add it, and wait for the container
 		owi := wsk8s.GetOWIFromObject(&newPod.ObjectMeta)
-		d.ctxs[workspaceInstanceID] = &workspaceState{
+		s := &workspaceState{
 			WorkspaceAdded: false,
 			Workspace: &Workspace{
 				InstanceID:  workspaceInstanceID,
 				WorkspaceID: workspaceID,
 				Pod:         newPod,
 			},
 		}
+		d.ctxs[workspaceInstanceID] = s
 
-		// Important!!!!: ideally this timeout must be equal to ws-manager https://github.com/gitpod-io/gitpod/blob/main/components/ws-manager/pkg/manager/manager.go#L171
-		waitForPodCtx, cancel := context.WithTimeout(context.Background(), 10*time.Minute)
 		containerCtx, containerCtxCancel := context.WithCancel(context.Background())
 		containerCtx = context.WithValue(containerCtx, contextDispatch, d)
+		containerCtx = context.WithValue(containerCtx, contextDispatchWaitGroup, &s.HandlerWaitGroup)
+		// Important!!!!: ideally this timeout must be equal to ws-manager https://github.com/gitpod-io/gitpod/blob/main/components/ws-manager/pkg/manager/manager.go#L171
+		waitForPodCtx, cancel := context.WithTimeout(containerCtx, 10*time.Minute)
 		go func() {
 			containerID, err := d.Runtime.WaitForContainer(waitForPodCtx, workspaceInstanceID)
 			if err != nil && err != context.Canceled {
@@ -217,12 +260,19 @@ func (d *Dispatch) handlePodUpdate(oldPod, newPod *corev1.Pod) {
 				d.mu.Unlock()
 				return
 			}
+			// Only register with the WaitGroup _after_ acquiring the lock to avoid DeadLocks
+			s.HandlerWaitGroup.Add(1)
+			defer s.HandlerWaitGroup.Done()
+
 			s.Context = containerCtx
-			s.Cancel = containerCtxCancel
+			s.Cancel = sync.OnceFunc(containerCtxCancel)
 			s.Workspace.ContainerID = containerID
 
 			for _, l := range d.Listener {
+				s.HandlerWaitGroup.Add(1)
 				go func(listener Listener) {
+					defer s.HandlerWaitGroup.Done()
+
 					err := listener.WorkspaceAdded(containerCtx, s.Workspace)
 					if err != nil {
 						log.WithError(err).WithFields(owi).Error("dispatch listener failed")
@@ -259,20 +309,25 @@ func (d *Dispatch) handlePodUpdate(oldPod, newPod *corev1.Pod) {
 			continue
 		}
 
+		state.HandlerWaitGroup.Add(1)
 		go func() {
+			defer state.HandlerWaitGroup.Done()
+
 			err := lu.WorkspaceUpdated(state.Context, state.Workspace)
 			if err != nil {
 				log.WithError(err).WithFields(wsk8s.GetOWIFromObject(&oldPod.ObjectMeta)).Error("dispatch listener failed")
 			}
 		}()
 	}
+	log.WithField("instanceID", workspaceInstanceID).Debugf("POD UPDATE DONE: %s", workspaceInstanceID)
 }
 
 func (d *Dispatch) handlePodDeleted(pod *corev1.Pod) {
 	instanceID, ok := pod.Labels[wsk8s.WorkspaceIDLabel]
 	if !ok {
 		return
 	}
+	log.WithField("instanceID", instanceID).Debugf("POD DELETED: %s", instanceID)
 
 	d.mu.Lock()
 	defer d.mu.Unlock()
@@ -286,4 +341,5 @@ func (d *Dispatch) handlePodDeleted(pod *corev1.Pod) {
 		state.Cancel()
 	}
 	delete(d.ctxs, instanceID)
+	log.WithField("instanceID", instanceID).Debugf("POD DELETED DONE: %s", instanceID)
 }

components/ws-daemon/pkg/netlimit/netlimit.go

@@ -128,8 +128,11 @@ func (c *ConnLimiter) limitWorkspace(ctx context.Context, ws *dispatch.Workspace
 		return fmt.Errorf("no dispatch available")
 	}
 
-	pid, err := disp.Runtime.ContainerPID(context.Background(), ws.ContainerID)
+	pid, err := disp.Runtime.ContainerPID(ctx, ws.ContainerID)
 	if err != nil {
+		if dispatch.IsCancelled(ctx) {
+			return nil
+		}
 		return fmt.Errorf("could not get pid for container %s of workspace %s", ws.ContainerID, ws.WorkspaceID)
 	}
 
@@ -141,12 +144,18 @@ func (c *ConnLimiter) limitWorkspace(ctx context.Context, ws *dispatch.Workspace
 		}
 	}, nsinsider.EnterMountNS(false), nsinsider.EnterNetNS(true))
 	if err != nil {
+		if dispatch.IsCancelled(ctx) {
+			return nil
+		}
 		log.WithError(err).WithFields(ws.OWI()).Error("cannot enable connection limiting")
 		return err
 	}
 	c.limited[ws.InstanceID] = struct{}{}
 
+	dispatch.GetDispatchWaitGroup(ctx).Add(1)
 	go func(*dispatch.Workspace) {
+		defer dispatch.GetDispatchWaitGroup(ctx).Done()
+
 		ticker := time.NewTicker(30 * time.Second)
 		defer ticker.Stop()