Update examples/stream.go

1 files changed, 29 insertions, 3 deletions

diff --git a/examples/stream.go b/examples/stream.go
index 1a6e2e4..56651c9 100644
--- a/examples/stream.go
+++ b/examples/stream.go
@@ -5,6 +5,7 @@
 package main
 
 import (
+	"context"
 	"crypto/tls"
 	"crypto/x509/pkix"
 	"fmt"
@@ -35,10 +36,35 @@ func main() {
 	}
 }
 
+// stream writes an infinite stream to w.
 func stream(w *gemini.ResponseWriter, r *gemini.Request) {
+	ch := make(chan string)
+	ctx, cancel := context.WithCancel(context.Background())
+
+	go func(ctx context.Context) {
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			default:
+				ch <- fmt.Sprint(time.Now().UTC())
+			}
+			time.Sleep(time.Second)
+		}
+		// Close channel when finished.
+		// In this example this will never be reached.
+		close(ch)
+	}(ctx)
+
 	for {
-		fmt.Fprintln(w, time.Now().UTC())
-		w.Flush()
-		time.Sleep(time.Second)
+		s, ok := <-ch
+		if !ok {
+			break
+		}
+		fmt.Fprintln(w, s)
+		if err := w.Flush(); err != nil {
+			cancel()
+			return
+		}
 	}
 }