Properly close grpc channels in uast extractor

Signed-off-by: Vadim Markovtsev <[email protected]>

Author: vmarkovtsev

internal/core/pipeline.go

@@ -128,6 +128,17 @@ type FeaturedPipelineItem interface {
 	Features() []string
 }
 
+// DisposablePipelineItem enables resources cleanup after finishing running the pipeline.
+type DisposablePipelineItem interface {
+	PipelineItem
+	// Dispose frees any previously allocated unmanaged resources. No Consume() calls are possible
+	// afterwards. The item needs to be Initialize()-d again.
+	// This method is invoked once for each item in the pipeline, **in a single forked instance**.
+	// Thus it is the responsibility of the item's programmer to deal with forks and merges, if
+	// necessary.
+	Dispose()
+}
+
 // LeafPipelineItem corresponds to the top level pipeline items which produce the end results.
 type LeafPipelineItem interface {
 	PipelineItem
@@ -848,6 +859,9 @@ func (pipeline *Pipeline) Run(commits []*object.Commit) (map[LeafPipelineItem]in
 	result := map[LeafPipelineItem]interface{}{}
 	if !pipeline.DryRun {
 		for index, item := range getMasterBranch(branches) {
+			if casted, ok := item.(DisposablePipelineItem); ok {
+				casted.Dispose()
+			}
 			if casted, ok := item.(LeafPipelineItem); ok {
 				result[pipeline.items[index].(LeafPipelineItem)] = casted.Finalize()
 			}

internal/core/pipeline_test.go

@@ -21,6 +21,7 @@ import (
 type testPipelineItem struct {
 	Initialized      bool
 	DepsConsumed     bool
+	Disposed         bool
 	Forked           bool
 	Merged           *bool
 	CommitMatches    bool
@@ -118,6 +119,10 @@ func (item *testPipelineItem) Consume(deps map[string]interface{}) (map[string]i
 	return map[string]interface{}{"test": item}, nil
 }
 
+func (item *testPipelineItem) Dispose() {
+	item.Disposed = true
+}
+
 func (item *testPipelineItem) Fork(n int) []PipelineItem {
 	result := make([]PipelineItem, n)
 	for i := 0; i < n; i++ {
@@ -304,6 +309,7 @@ func TestPipelineRun(t *testing.T) {
 		assert.True(t, val >= 0, key)
 	}
 	assert.True(t, item.DepsConsumed)
+	assert.True(t, item.Disposed)
 	assert.True(t, item.CommitMatches)
 	assert.True(t, item.IndexMatches)
 	assert.Equal(t, 1, *item.MergeState)

internal/plumbing/uast/uast.go

@@ -294,6 +294,13 @@ func (exr *Extractor) Consume(deps map[string]interface{}) (map[string]interface
 	return map[string]interface{}{DependencyUasts: uasts}, nil
 }
 
+// Dispose closes the open GRPC channels.
+func (exr *Extractor) Dispose() {
+	for _, client := range exr.clients {
+		client.Close()
+	}
+}
+
 // Fork clones this PipelineItem.
 func (exr *Extractor) Fork(n int) []core.PipelineItem {
 	return core.ForkSamePipelineItem(exr, n)

internal/plumbing/uast/uast_test.go

@@ -43,6 +43,7 @@ func AddHash(t *testing.T, cache map[plumbing.Hash]*items.CachedBlob, hash strin
 
 func TestUASTExtractorMeta(t *testing.T) {
 	exr := fixtureUASTExtractor()
+	defer exr.Dispose()
 	assert.Equal(t, exr.Name(), "UAST")
 	assert.Equal(t, len(exr.Provides()), 1)
 	assert.Equal(t, exr.Provides()[0], DependencyUasts)
@@ -63,14 +64,15 @@ func TestUASTExtractorMeta(t *testing.T) {
 
 func TestUASTExtractorConfiguration(t *testing.T) {
 	exr := fixtureUASTExtractor()
+	defer exr.Dispose()
 	facts := map[string]interface{}{}
-	exr.Configure(facts)
+	assert.Nil(t, exr.Configure(facts))
 	facts[ConfigUASTEndpoint] = "localhost:9432"
 	facts[ConfigUASTTimeout] = 15
 	facts[ConfigUASTPoolSize] = 7
 	facts[ConfigUASTFailOnErrors] = true
 	facts[ConfigUASTIgnoreMissingDrivers] = []string{"test"}
-	exr.Configure(facts)
+	assert.Nil(t, exr.Configure(facts))
 	assert.Equal(t, exr.Endpoint, facts[ConfigUASTEndpoint])
 	assert.NotNil(t, exr.Context)
 	assert.Equal(t, exr.PoolSize, facts[ConfigUASTPoolSize])
@@ -95,6 +97,7 @@ func TestUASTExtractorNoBabelfish(t *testing.T) {
 
 func TestUASTExtractorConsume(t *testing.T) {
 	exr := fixtureUASTExtractor()
+	defer exr.Dispose()
 	changes := make(object.Changes, 4)
 	// 2b1ed978194a94edeabbca6de7ff3b5771d4d665
 	treeFrom, _ := test.Repository.TreeObject(plumbing.NewHash(
@@ -211,6 +214,7 @@ func TestUASTExtractorConsume(t *testing.T) {
 
 func TestUASTExtractorFork(t *testing.T) {
 	exr1 := fixtureUASTExtractor()
+	defer exr1.Dispose()
 	clones := exr1.Fork(1)
 	assert.Len(t, clones, 1)
 	exr2 := clones[0].(*Extractor)