Add context cancellation

five.diff

diff --git a/main.go b/main.go
index 606c863..08d7580 100644
--- a/main.go
+++ b/main.go
@@ -1,37 +1,58 @@
 package main
 
 import (
+	"context"
 	"log"
 	"time"
 )
 
-func producer(strings []string) (<-chan string, error) {
+func producer(ctx context.Context, strings []string) (<-chan string, error) {
 	outChannel := make(chan string)
 
 	go func() {
 		defer close(outChannel)
 		for _, s := range strings {
 			time.Sleep(time.Second * 3)
-			outChannel <- s
+			select {
+			case <-ctx.Done():
+				return
+			case outChannel <- s:
+			}
 		}
 	}()
 
 	return outChannel, nil
 }
 
-func sink(values <-chan string) {
-	for value := range values {
-		log.Println(value)
+func sink(ctx context.Context, values <-chan string) {
+	for {
+		select {
+		case <-ctx.Done():
+			log.Print(ctx.Err().Error())
+			return
+		case val, ok := <-values:
+			if ok {
+				log.Printf("sink: %s", val)
+			}
+		}
 	}
 }
 
 func main() {
 	source := []string{"foo", "bar", "bax"}
 
-	outputChannel, err := producer(source)
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	go func() {
+		time.Sleep(time.Second * 5)
+		cancel()
+	}()
+
+	outputChannel, err := producer(ctx, source)
 	if err != nil {
 		log.Fatal(err)
 	}
 
-	sink(outputChannel)
-}
\ No newline at end of file
+	sink(ctx, outputChannel)
+}