diff --git a/input_spliter.go b/input_spliter.go index da1ff5c..11516bc 100644 --- a/input_spliter.go +++ b/input_spliter.go @@ -24,13 +24,16 @@ func split(ctx context.Context, r io.Reader) (<-chan string, <-chan error) { for sc.Scan() { select { case <-ctx.Done(): - errc <- ctx.Err() return default: l := sc.Text() if isRootBlockBeginning(l) { if len(block) != 0 { - blockc <- block + select { + case <-ctx.Done(): + return + case blockc <- block: + } } block = "" } @@ -41,8 +44,12 @@ func split(ctx context.Context, r io.Reader) (<-chan string, <-chan error) { errc <- err return } - blockc <- block // 最後のRootブロック送出 - return + select { + case <-ctx.Done(): + return + case blockc <- block: // 最後のRootブロック送出 + return + } }() return blockc, errc diff --git a/pipeline_tree_grower.go b/pipeline_tree_grower.go index 4c7808d..3e7c829 100644 --- a/pipeline_tree_grower.go +++ b/pipeline_tree_grower.go @@ -57,7 +57,11 @@ func (dg *defaultGrowerPipeline) worker(ctx context.Context, wg *sync.WaitGroup, errc <- err return } - nodes <- root + select { + case <-ctx.Done(): + return + case nodes <- root: + } } } } diff --git a/root_generator.go b/root_generator.go index 5ea2e25..e7ebfa5 100644 --- a/root_generator.go +++ b/root_generator.go @@ -134,8 +134,11 @@ func (rg *rootGeneratorPipeline) worker(ctx context.Context, wg *sync.WaitGroup, errc <- err return } - - rootc <- root + select { + case <-ctx.Done(): + return + case rootc <- root: + } } } } diff --git a/tree_handler_output_test.go b/tree_handler_output_test.go index f9a2e69..dd1b733 100644 --- a/tree_handler_output_test.go +++ b/tree_handler_output_test.go @@ -10,6 +10,7 @@ import ( "os" "strings" "testing" + "time" "github.com/ddddddO/gtree" tu "github.com/ddddddO/gtree/testutil" @@ -1192,3 +1193,15 @@ children: }) } } + +func TestOutput_detecting_goroutinelerk(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), time.Duration(500*time.Millisecond)) + defer cancel() + w := &strings.Builder{} + r := strings.NewReader(tu.TwentyThousandRoots) + if gotErr := gtree.Output(w, r, gtree.WithMassive(ctx)); gotErr != nil { + if gotErr != context.DeadlineExceeded { + t.Errorf("\ngotErr: \n%v\nwantErr: \n%v", gotErr, context.DeadlineExceeded) + } + } +}