From 90b989a7186ffdb83d3e52110199fcbb5295d665 Mon Sep 17 00:00:00 2001
From: Juan Batiz-Benet <juan@benet.ai>
Date: Sat, 25 Oct 2014 05:21:51 -0700
Subject: [PATCH] AddCloserChild + net ctxcloser

---
 net/net.go               | 36 ++++++++++--------------------------
 util/ctxcloser/closer.go | 24 ++++++++++++++++++++++++
 2 files changed, 34 insertions(+), 26 deletions(-)

diff --git a/net/net.go b/net/net.go
index 8d6d34d30..a6154a780 100644
--- a/net/net.go
+++ b/net/net.go
@@ -1,12 +1,11 @@
 package net
 
 import (
-	"errors"
-
 	msg "github.com/jbenet/go-ipfs/net/message"
 	mux "github.com/jbenet/go-ipfs/net/mux"
 	swarm "github.com/jbenet/go-ipfs/net/swarm"
 	peer "github.com/jbenet/go-ipfs/peer"
+	ctxc "github.com/jbenet/go-ipfs/util/ctxcloser"
 
 	context "github.com/jbenet/go-ipfs/Godeps/_workspace/src/code.google.com/p/go.net/context"
 )
@@ -23,31 +22,30 @@ type IpfsNetwork struct {
 	// peer connection multiplexing
 	swarm *swarm.Swarm
 
-	// network context
-	ctx    context.Context
-	cancel context.CancelFunc
+	// network context closer
+	ctxc.ContextCloser
 }
 
 // NewIpfsNetwork is the structure that implements the network interface
 func NewIpfsNetwork(ctx context.Context, local peer.Peer,
 	peers peer.Peerstore, pmap *mux.ProtocolMap) (*IpfsNetwork, error) {
 
-	ctx, cancel := context.WithCancel(ctx)
-
 	in := &IpfsNetwork{
-		local:  local,
-		muxer:  mux.NewMuxer(ctx, *pmap),
-		ctx:    ctx,
-		cancel: cancel,
+		local:         local,
+		muxer:         mux.NewMuxer(ctx, *pmap),
+		ContextCloser: ctxc.NewContextCloser(ctx, nil),
 	}
 
 	var err error
 	in.swarm, err = swarm.NewSwarm(ctx, local, peers)
 	if err != nil {
-		cancel()
+		in.Close()
 		return nil, err
 	}
 
+	in.AddCloserChild(in.swarm)
+	in.AddCloserChild(in.muxer)
+
 	// remember to wire components together.
 	in.muxer.Pipe.ConnectTo(in.swarm.Pipe)
 
@@ -89,20 +87,6 @@ func (n *IpfsNetwork) SendMessage(m msg.NetMessage) error {
 	return nil
 }
 
-// Close terminates all network operation
-func (n *IpfsNetwork) Close() error {
-	if n.cancel == nil {
-		return errors.New("Network already closed.")
-	}
-
-	n.swarm.Close()
-	n.muxer.Close()
-
-	n.cancel()
-	n.cancel = nil
-	return nil
-}
-
 // GetPeerList returns the networks list of connected peers
 func (n *IpfsNetwork) GetPeerList() []peer.Peer {
 	return n.swarm.GetPeerList()
diff --git a/util/ctxcloser/closer.go b/util/ctxcloser/closer.go
index 348951ad6..ca80368ea 100644
--- a/util/ctxcloser/closer.go
+++ b/util/ctxcloser/closer.go
@@ -50,10 +50,25 @@ type ContextCloser interface {
 	// Children is a sync.Waitgroup for all children goroutines that should
 	// shut down completely before this service is said to be "closed".
 	// Follows the semantics of WaitGroup:
+	//
 	//  Children().Add(1) // add one more dependent child
 	//  Children().Done() // child signals it is done
+	//
 	Children() *sync.WaitGroup
 
+	// AddCloserChild registers a dependent ContextCloser child. The child will
+	// be closed when this parent is closed, and waited upon to finish. It is
+	// the functional equivalent of the following:
+	//
+	//  go func(parent, child ContextCloser) {
+	//  	parent.Children().Add(1) // add one more dependent child
+	//  	<-parent.Closing()       // wait until parent is closing
+	//  	child.Close()            // signal child to close
+	//  	parent.Children().Done() // child signals it is done
+	//	}(a, b)
+	//
+	AddCloserChild(c ContextCloser)
+
 	// Close is a method to call when you wish to stop this ContextCloser
 	Close() error
 
@@ -117,6 +132,15 @@ func (c *contextCloser) Children() *sync.WaitGroup {
 	return &c.children
 }
 
+func (c *contextCloser) AddCloserChild(child ContextCloser) {
+	c.children.Add(1)
+	go func(parent, child ContextCloser) {
+		<-parent.Closing()       // wait until parent is closing
+		child.Close()            // signal child to close
+		parent.Children().Done() // child signals it is done
+	}(c, child)
+}
+
 // Close is the external close function. it's a wrapper around internalClose
 // that waits on Closed()
 func (c *contextCloser) Close() error {