From 03ee057a6eeb9a986d6a41ad05503b089118dc17 Mon Sep 17 00:00:00 2001
From: Simon Richardson <stickupkid@gmail.com>
Date: Thu, 30 Mar 2023 12:00:50 +0100
Subject: [PATCH] Catacomb context support

Adds context support to the catacomb type. In reality this just calls
the tomb context method. By exposing the context method we can then tie
a context to the lifecycle of a catacomb.

The only difference between a tomb and a catacomb in relation to a
context is that a catacomb doesn't expose the concept of a parent
context. The tomb type allows you to create a tomb from a context
(WithContext), which will be the root context for the tomb. Instead a
catacomb will currently fallback to a context.Background. We could
provide a WithContext for a catacomb, I'm unsure if that functionality
is required right now?

---

The tests were lifted and modified from the tomb test package, which
ensures that we don't deviate from the tomb package in the future.
---
 catacomb/catacomb.go      | 19 ++++++++--
 catacomb/catacomb_test.go | 78 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 93 insertions(+), 4 deletions(-)

diff --git a/catacomb/catacomb.go b/catacomb/catacomb.go
index 67ff35a..0d333c7 100644
--- a/catacomb/catacomb.go
+++ b/catacomb/catacomb.go
@@ -4,6 +4,7 @@
 package catacomb
 
 import (
+	"context"
 	"fmt"
 	"runtime/debug"
 	"strings"
@@ -211,12 +212,22 @@ func (catacomb *Catacomb) Err() error {
 	return catacomb.tomb.Err()
 }
 
+// Context returns a context that is a copy of the provided parent context with
+// a replaced Done channel that is closed when either the catacomb is dying or
+// the parent is cancelled.
+//
+// If parent is nil, it defaults to the empty background context.
+func (catacomb *Catacomb) Context(parent context.Context) context.Context {
+	return catacomb.tomb.Context(parent)
+}
+
 // Kill kills the Catacomb's internal tomb with the supplied error, or one
 // derived from it.
-//  * if it's caused by this catacomb's ErrDying, it passes on tomb.ErrDying.
-//  * if it's tomb.ErrDying, or caused by another catacomb's ErrDying, it passes
-//    on a new error complaining about the misuse.
-//  * all other errors are passed on unmodified.
+//   - if it's caused by this catacomb's ErrDying, it passes on tomb.ErrDying.
+//   - if it's tomb.ErrDying, or caused by another catacomb's ErrDying, it passes
+//     on a new error complaining about the misuse.
+//   - all other errors are passed on unmodified.
+//
 // It's always safe to call Kill, but errors passed to Kill after the catacomb
 // is dead will be ignored.
 func (catacomb *Catacomb) Kill(err error) {
diff --git a/catacomb/catacomb_test.go b/catacomb/catacomb_test.go
index eae770a..6db34d5 100644
--- a/catacomb/catacomb_test.go
+++ b/catacomb/catacomb_test.go
@@ -4,6 +4,7 @@
 package catacomb_test
 
 import (
+	"context"
 	"sync"
 
 	"github.com/juju/errors"
@@ -481,3 +482,80 @@ func checkInvalid(c *gc.C, plan catacomb.Plan, match string) {
 	check(plan.Validate())
 	check(catacomb.Invoke(plan))
 }
+
+type CatacombContextSuite struct {
+	testing.IsolationSuite
+}
+
+var _ = gc.Suite(&CatacombContextSuite{})
+
+func (CatacombContextSuite) TestContextNoParent(c *gc.C) {
+	var site catacomb.Catacomb
+	err := catacomb.Invoke(catacomb.Plan{
+		Site: &site,
+		Work: func() error { return nil },
+	})
+	c.Check(err, jc.ErrorIsNil)
+
+	parent2, cancel2 := context.WithCancel(context.WithValue(context.Background(), "parent", "parent2"))
+	child2 := site.Context(parent2)
+
+	if site.Context(parent2) != child2 {
+		c.Fatalf("Context returned different context for same parent")
+	}
+	if child2.Value("parent") != "parent2" {
+		c.Fatalf("Child context didn't inherit its parent's properties")
+	}
+	select {
+	case <-child2.Done():
+		c.Fatalf("Tomb's child context was born dead")
+	default:
+	}
+
+	cancel2()
+
+	select {
+	case <-child2.Done():
+	default:
+		c.Fatalf("Tomb's child context didn't die after parent was canceled")
+	}
+
+	parent3 := context.WithValue(context.Background(), "parent", "parent3")
+	child3 := site.Context(parent3)
+
+	if child3.Value("parent") != "parent3" {
+		c.Fatalf("Child context didn't inherit its parent's properties")
+	}
+	select {
+	case <-child3.Done():
+		c.Fatalf("Tomb's child context was born dead")
+	default:
+	}
+
+	site.Kill(nil)
+
+	if site.Context(parent3) == child3 {
+		c.Fatalf("Tomb is dead and shouldn't be tracking children anymore")
+	}
+	select {
+	case <-child3.Done():
+	default:
+		c.Fatalf("Child context didn't die after tomb's death")
+	}
+
+	parent4 := context.WithValue(context.Background(), "parent", "parent4")
+	child4 := site.Context(parent4)
+
+	select {
+	case <-child4.Done():
+	default:
+		c.Fatalf("Child context should be born canceled")
+	}
+
+	childnil := site.Context(nil)
+	select {
+	case <-childnil.Done():
+	default:
+		c.Fatalf("Child context should be born canceled")
+	}
+}