massive with context (#189)

cmd/gtree/main.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"log"
@@ -193,7 +194,8 @@ func actionOutput(c *cli.Context) error {
 	}
 	options := []gtree.Option{oi, oo}
 	if c.Bool("massive") {
-		options = append(options, gtree.WithMassive())
+		// TODO: ほぼいらないけど、タイムアウト値をフラグで取ってcontext.WithTimeout渡すようにするでもいいかも
+		options = append(options, gtree.WithMassive(context.Background()))
 	}
 
 	markdownPath := c.Path("file")
@@ -243,7 +245,7 @@ func actionMkdir(c *cli.Context) error {
 	oe := gtree.WithFileExtensions(c.StringSlice("extension"))
 	options := []gtree.Option{oi, oe}
 	if c.Bool("massive") {
-		options = append(options, gtree.WithMassive())
+		options = append(options, gtree.WithMassive(context.Background()))
 	}
 
 	if c.Bool("dry-run") {

config.go

@@ -1,11 +1,14 @@
 package gtree
 
+import "context"
+
 type config struct {
 	lastNodeFormat        branchFormat
 	intermedialNodeFormat branchFormat
 
 	space          spaceType
 	massive        bool
+	ctx            context.Context
 	encode         encode
 	dryrun         bool
 	fileExtensions []string
@@ -68,9 +71,14 @@ func WithBranchFormatLastNode(directly, indirectly string) Option {
 }
 
 // WithMassive returns function for large amount roots.
-func WithMassive() Option {
+func WithMassive(ctx context.Context) Option {
 	return func(c *config) {
 		c.massive = true
+
+		if ctx == nil {
+			ctx = context.Background()
+		}
+		c.ctx = ctx
 	}
 }
 

input_spliter.go

@@ -20,30 +20,29 @@ func split(ctx context.Context, r io.Reader) (<-chan string, <-chan error) {
 			close(errc)
 		}()
 
-		for {
+		block := ""
+		for sc.Scan() {
 			select {
 			case <-ctx.Done():
+				errc <- ctx.Err()
 				return
 			default:
-				block := ""
-				for sc.Scan() {
-					l := sc.Text()
-					if isRootBlockBeginning(l) {
-						if len(block) != 0 {
-							blockc <- block
-						}
-						block = ""
+				l := sc.Text()
+				if isRootBlockBeginning(l) {
+					if len(block) != 0 {
+						blockc <- block
 					}
-					block += fmt.Sprintln(l)
+					block = ""
 				}
-				if err := sc.Err(); err != nil {
-					errc <- err
-					return
-				}
-				blockc <- block // 最後のRootブロック送出
-				return
+				block += fmt.Sprintln(l)
 			}
 		}
+		if err := sc.Err(); err != nil {
+			errc <- err
+			return
+		}
+		blockc <- block // 最後のRootブロック送出
+		return
 	}()
 
 	return blockc, errc

pipeline_tree.go

@@ -54,18 +54,18 @@ func newTreePipeline(conf *config) iTree {
 }
 
 func (t *treePipeline) output(w io.Writer, r io.Reader, conf *config) error {
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx, cancel := context.WithCancel(conf.ctx)
 	defer cancel()
 
 	splitStream, errcsl := split(ctx, r)
 	rootStream, errcr := newRootGeneratorPipeline(conf.space).generate(ctx, splitStream)
 	growStream, errcg := t.grow(ctx, rootStream)
 	errcs := t.spread(ctx, w, growStream)
-	return t.handlePipelineErr(errcsl, errcr, errcg, errcs)
+	return t.handlePipelineErr(ctx, errcsl, errcr, errcg, errcs)
 }
 
 func (t *treePipeline) outputProgrammably(w io.Writer, root *Node, conf *config) error {
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx, cancel := context.WithCancel(conf.ctx)
 	defer cancel()
 
 	rootStream := make(chan *Node)
@@ -75,22 +75,22 @@ func (t *treePipeline) outputProgrammably(w io.Writer, root *Node, conf *config)
 	}()
 	growStream, errcg := t.grow(ctx, rootStream)
 	errcs := t.spread(ctx, w, growStream)
-	return t.handlePipelineErr(errcg, errcs)
+	return t.handlePipelineErr(ctx, errcg, errcs)
 }
 
 func (t *treePipeline) makedir(r io.Reader, conf *config) error {
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx, cancel := context.WithCancel(conf.ctx)
 	defer cancel()
 
 	splitStream, errcsl := split(ctx, r)
 	rootStream, errcr := newRootGeneratorPipeline(conf.space).generate(ctx, splitStream)
 	growStream, errcg := t.grow(ctx, rootStream)
 	errcm := t.mkdir(ctx, growStream)
-	return t.handlePipelineErr(errcsl, errcr, errcg, errcm)
+	return t.handlePipelineErr(ctx, errcsl, errcr, errcg, errcm)
 }
 
 func (t *treePipeline) makedirProgrammably(root *Node, conf *config) error {
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx, cancel := context.WithCancel(conf.ctx)
 	defer cancel()
 
 	rootStream := make(chan *Node)
@@ -104,11 +104,11 @@ func (t *treePipeline) makedirProgrammably(root *Node, conf *config) error {
 	if conf.dryrun {
 		// when detected no invalid node name, output tree.
 		errcs := t.spread(ctx, color.Output, growStream)
-		return t.handlePipelineErr(errcg, errcs)
+		return t.handlePipelineErr(ctx, errcg, errcs)
 	}
 	// when detected no invalid node name, no output tree.
 	errcm := t.mkdir(ctx, growStream)
-	return t.handlePipelineErr(errcg, errcm)
+	return t.handlePipelineErr(ctx, errcg, errcm)
 }
 
 // 関心事は各ノードの枝の形成
@@ -141,17 +141,22 @@ func (t *treePipeline) mkdir(ctx context.Context, roots <-chan *Node) <-chan err
 }
 
 // パイプラインの全ステージで最初のエラーを返却
-func (*treePipeline) handlePipelineErr(echs ...<-chan error) error {
-	eg, _ := errgroup.WithContext(context.TODO())
+func (*treePipeline) handlePipelineErr(ctx context.Context, echs ...<-chan error) error {
+	eg, ectx := errgroup.WithContext(ctx)
 	for i := range echs {
 		i := i
 		eg.Go(func() error {
-			for e := range echs[i] {
-				if e != nil {
-					return e
+			for {
+				select {
+				case err := <-echs[i]:
+					if err != nil {
+						return err
+					}
+					return nil
+				case <-ectx.Done():
+					return ectx.Err()
 				}
 			}
-			return nil
 		})
 	}
 	return eg.Wait()

pipeline_tree_spreader.go

@@ -74,7 +74,9 @@ func (ds *defaultSpreaderPipeline) worker(ctx context.Context, wg *sync.WaitGrou
 			_, err := bw.WriteString(ret)
 			ds.Unlock()
 
-			errc <- err
+			if err != nil {
+				errc <- err
+			}
 		}
 	}
 }

testutil/benchmark_util.go

@@ -1,14 +1,15 @@
 package testutil
 
 import (
+	"context"
 	"strings"
 	"testing"
 
 	"github.com/ddddddO/gtree"
 )
 
 func BaseBenchmarkWithMassive(roots string, b *testing.B) {
-	BaseBenchmark(roots, b, gtree.WithMassive())
+	BaseBenchmark(roots, b, gtree.WithMassive(context.Background()))
 }
 
 func BaseBenchmark(roots string, b *testing.B, options ...gtree.Option) {

tree_handler_mkdir_test.go

@@ -1,6 +1,7 @@
 package gtree_test
 
 import (
+	"context"
 	"errors"
 	"strings"
 	"testing"
@@ -48,7 +49,7 @@ func TestMkdir(t *testing.T) {
 	- bb
 		- lll
 	-ff`)),
-				options: []gtree.Option{gtree.WithMassive()},
+				options: []gtree.Option{gtree.WithMassive(context.Background())},
 			},
 			wantErr: gtree.ErrExistPath,
 		},
@@ -131,7 +132,7 @@ func TestMkdir(t *testing.T) {
 		- lll`)),
 				options: []gtree.Option{
 					gtree.WithFileExtensions([]string{".go"}),
-					gtree.WithMassive(),
+					gtree.WithMassive(context.Background()),
 				},
 			},
 			wantErr: nil,

tree_handler_output_test.go

@@ -3,6 +3,7 @@ package gtree_test
 import (
 	"bufio"
 	"bytes"
+	"context"
 	"errors"
 	"fmt"
 	"io"
@@ -348,7 +349,7 @@ a
 	- e
 		- o
 	- g`)),
-				options: []gtree.Option{gtree.WithMassive()},
+				options: []gtree.Option{gtree.WithMassive(context.Background())},
 			},
 			out: out{
 				output: strings.TrimPrefix(`
@@ -417,7 +418,7 @@ a
 				input: strings.NewReader(strings.TrimSpace(fmt.Sprintf(`
 - a
 	- %s`, strings.Repeat("a", 64*1024)))),
-				options: []gtree.Option{gtree.WithMassive()},
+				options: []gtree.Option{gtree.WithMassive(context.Background())},
 			},
 			out: out{
 				output: "",
@@ -560,7 +561,7 @@ a	prev tab
 	- b
 		- c`)),
 				options: []gtree.Option{
-					gtree.WithMassive(),
+					gtree.WithMassive(context.Background()),
 				},
 			},
 			out: out{
@@ -583,7 +584,7 @@ a
 		- c
 	- y`)),
 				options: []gtree.Option{
-					gtree.WithMassive(),
+					gtree.WithMassive(context.Background()),
 					gtree.WithDryRun(),
 					gtree.WithFileExtensions([]string{"c"}),
 				},
@@ -610,7 +611,7 @@ a
 	- b
 		- c`)),
 				options: []gtree.Option{
-					gtree.WithMassive(),
+					gtree.WithMassive(context.Background()),
 					gtree.WithEncodeJSON(),
 				},
 			},
@@ -628,7 +629,7 @@ a
 	- b
 		- c`)),
 				options: []gtree.Option{
-					gtree.WithMassive(),
+					gtree.WithMassive(context.Background()),
 					gtree.WithEncodeYAML(),
 				},
 			},
@@ -653,7 +654,7 @@ children:
 	- b
 		- c`)),
 				options: []gtree.Option{
-					gtree.WithMassive(),
+					gtree.WithMassive(context.Background()),
 					gtree.WithEncodeTOML(),
 				},
 			},

tree_handler_programmably_mkdir_test.go

@@ -1,6 +1,7 @@
 package gtree_test
 
 import (
+	"context"
 	"fmt"
 	"testing"
 
@@ -22,7 +23,7 @@ func TestMkdirProgrammably(t *testing.T) {
 		{
 			name:    "case(succeeded/massive)",
 			root:    tu.Prepare_a(),
-			options: []gtree.Option{gtree.WithMassive()},
+			options: []gtree.Option{gtree.WithMassive(context.Background())},
 		},
 		{
 			name:    "case(not root)",

tree_handler_programmably_output_test.go

@@ -2,6 +2,7 @@ package gtree_test
 
 import (
 	"bytes"
+	"context"
 	"fmt"
 	"os"
 	"strings"
@@ -59,7 +60,7 @@ root
 		{
 			name:    "case(succeeded/massive)",
 			root:    tu.Prepare(),
-			options: []gtree.Option{gtree.WithMassive()},
+			options: []gtree.Option{gtree.WithMassive(context.Background())},
 			want: strings.TrimPrefix(`
 root
 └── child 1