From 6eea8f3f627672e05afc340902b3bd2931d596a2 Mon Sep 17 00:00:00 2001
From: Derek McGowan <derek@mcg.dev>
Date: Thu, 16 Sep 2021 17:56:28 -0700
Subject: [PATCH] Add shutdown package

Allows shutdown to handle callbacks with similar behavior as context
cancel

Signed-off-by: Derek McGowan <derek@mcg.dev>
---
 pkg/shutdown/shutdown.go | 108 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 108 insertions(+)
 create mode 100644 pkg/shutdown/shutdown.go

diff --git a/pkg/shutdown/shutdown.go b/pkg/shutdown/shutdown.go
new file mode 100644
index 000000000..982a95099
--- /dev/null
+++ b/pkg/shutdown/shutdown.go
@@ -0,0 +1,108 @@
+/*
+   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 shutdown
+
+import (
+	"context"
+	"errors"
+	"sync"
+	"time"
+
+	"golang.org/x/sync/errgroup"
+)
+
+// ErrShutdown is the error condition when a context has been fully shutdown
+var ErrShutdown = errors.New("shutdown")
+
+// Service is used to facilitate shutdown by through callback
+// registration and shutdown initiation
+type Service interface {
+	// Shutdown initiates shutdown
+	Shutdown()
+	// RegisterCallback registers functions to be called on shutdown and before
+	// the shutdown channel is closed. A callback error will propagate to the
+	// context error
+	RegisterCallback(func(context.Context) error)
+}
+
+// WithShutdown returns a context which is similar to a cancel context, but
+// with callbacks which can propagate to the context error. Unlike a cancel
+// context, the shutdown context cannot be canceled from the parent context.
+// However, future child contexes will be canceled upon shutdown.
+func WithShutdown(ctx context.Context) (context.Context, Service) {
+	ss := &shutdownService{
+		Context: ctx,
+		doneC:   make(chan struct{}),
+		timeout: 30 * time.Second,
+	}
+	return ss, ss
+}
+
+type shutdownService struct {
+	context.Context
+
+	mu         sync.Mutex
+	isShutdown bool
+	callbacks  []func(context.Context) error
+	doneC      chan struct{}
+	err        error
+	timeout    time.Duration
+}
+
+func (s *shutdownService) Shutdown() {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	if s.isShutdown {
+		return
+	}
+	s.isShutdown = true
+
+	go func(callbacks []func(context.Context) error) {
+		ctx, cancel := context.WithTimeout(context.Background(), s.timeout)
+		defer cancel()
+		grp, ctx := errgroup.WithContext(ctx)
+		for _, fn := range callbacks {
+			grp.Go(func() error { return fn(ctx) })
+		}
+		err := grp.Wait()
+		if err == nil {
+			err = ErrShutdown
+		}
+		s.mu.Lock()
+		s.err = err
+		close(s.doneC)
+		s.mu.Unlock()
+	}(s.callbacks)
+}
+
+func (s *shutdownService) Done() <-chan struct{} {
+	return s.doneC
+}
+
+func (s *shutdownService) Err() error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	return s.err
+}
+func (s *shutdownService) RegisterCallback(fn func(context.Context) error) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	if s.callbacks == nil {
+		s.callbacks = []func(context.Context) error{}
+	}
+	s.callbacks = append(s.callbacks, fn)
+}