Add cleanup package for context management during cleanup

dmcgowan · dmcgowan · commit f606c4eba738 · 2023-01-03T12:30:26.000-08:00
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 &lt;derek@mcg.dev&gt;
diff --git a/pkg/cleanup/context.go 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
@@ -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
+	}
+}
