From f606c4eba7386f64a8b4db9a5b2b88ef16e42657 Mon Sep 17 00:00:00 2001
From: Derek McGowan <derek@mcg.dev>
Date: Thu, 22 Dec 2022 16:44:08 -0800
Subject: [PATCH] Add cleanup package for context management during cleanup

Provides a couple helper functions that provide a background context for
running cleanup jobs while preserving the original context values.
The new contexts will not inherit the errors or cancellations.

Signed-off-by: Derek McGowan <derek@mcg.dev>
---
 pkg/cleanup/context.go      | 52 +++++++++++++++++++++++++++++++++
 pkg/cleanup/context_test.go | 58 +++++++++++++++++++++++++++++++++++++
 2 files changed, 110 insertions(+)
 create mode 100644 pkg/cleanup/context.go
 create mode 100644 pkg/cleanup/context_test.go

diff --git a/pkg/cleanup/context.go b/pkg/cleanup/context.go
new file mode 100644
index 000000000..62741c4ca
--- /dev/null
+++ b/pkg/cleanup/context.go
@@ -0,0 +1,52 @@
+/*
+   Copyright The containerd Authors.
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+*/
+
+// Package providing utilies to help cleanup
+package cleanup
+
+import (
+	"context"
+	"time"
+)
+
+type clearCancel struct {
+	context.Context
+}
+
+func (cc clearCancel) Deadline() (deadline time.Time, ok bool) {
+	return
+}
+
+func (cc clearCancel) Done() <-chan struct{} {
+	return nil
+}
+
+func (cc clearCancel) Err() error {
+	return nil
+}
+
+// Background creates a new context which clears out the parent errors
+func Background(ctx context.Context) context.Context {
+	return clearCancel{ctx}
+}
+
+// Do runs the provided function with a context in which the
+// errors are cleared out and will timeout after 10 seconds.
+func Do(ctx context.Context, do func(context.Context)) {
+	ctx, cancel := context.WithTimeout(clearCancel{ctx}, 10*time.Second)
+	do(ctx)
+	cancel()
+}
diff --git a/pkg/cleanup/context_test.go b/pkg/cleanup/context_test.go
new file mode 100644
index 000000000..ced1c3d41
--- /dev/null
+++ b/pkg/cleanup/context_test.go
@@ -0,0 +1,58 @@
+/*
+   Copyright The containerd Authors.
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+*/
+
+package cleanup
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestBackground(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	var k struct{}
+	v := "incontext"
+	ctx = context.WithValue(ctx, k, v)
+
+	assert.Nil(t, contextError(ctx))
+	assert.Equal(t, ctx.Value(k), v)
+
+	cancel()
+	assert.Error(t, contextError(ctx))
+	assert.Equal(t, ctx.Value(k), v)
+
+	// cleanup context should no longer be canceled
+	ctx = Background(ctx)
+	assert.Nil(t, contextError(ctx))
+	assert.Equal(t, ctx.Value(k), v)
+
+	// cleanup contexts can be rewrapped in cancel context
+	ctx, cancel = context.WithCancel(ctx)
+	cancel()
+	assert.Error(t, contextError(ctx))
+	assert.Equal(t, ctx.Value(k), v)
+}
+
+func contextError(ctx context.Context) error {
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	default:
+		return nil
+	}
+}