Add new Hydrate function wrapper Memoize to replace WithCache. Replace Mutexes with RWMutexes and update locking to use RLock where possible. Closes #498. Closes #499

Author: kaidaguerre

connection/connection_cache.go

@@ -33,6 +33,7 @@ func (c *ConnectionCache) SetWithTTL(ctx context.Context, key string, value inte
 	err := c.cache.Set(ctx,
 		key,
 		value,
+		// if ttl is zero there is no expiration
 		store.WithExpiration(ttl),
 		// put connection name in tags
 		store.WithTags([]string{c.connectionName}),

go.mod

@@ -36,6 +36,7 @@ require (
 	github.com/eko/gocache/v3 v3.1.2
 	github.com/fsnotify/fsnotify v1.6.0
 	github.com/hashicorp/go-getter v1.6.2
+	golang.org/x/exp v0.0.0-20221110155412-d0897a79cd37
 	golang.org/x/sync v0.1.0
 )
 
@@ -93,7 +94,6 @@ require (
 	go.opencensus.io v0.22.4 // indirect
 	go.opentelemetry.io/otel/exporters/otlp/internal/retry v1.7.0 // indirect
 	go.opentelemetry.io/proto/otlp v0.16.0 // indirect
-	golang.org/x/exp v0.0.0-20221110155412-d0897a79cd37 // indirect
 	golang.org/x/lint v0.0.0-20200302205851-738671d3881b // indirect
 	golang.org/x/mod v0.6.0 // indirect
 	golang.org/x/net v0.2.0 // indirect

grpc/proto/plugin.pb.go

@@ -9,20 +9,23 @@ import (
 DefaultConcurrencyConfig sets the default maximum number of concurrent [HydrateFunc] calls.
 
 Limit total concurrent hydrate calls:
-		DefaultConcurrency: &plugin.DefaultConcurrencyConfig{
-			TotalMaxConcurrency:   500,
-		}
+
+	DefaultConcurrency: &plugin.DefaultConcurrencyConfig{
+		TotalMaxConcurrency:   500,
+	}
 
 Limit concurrent hydrate calls to any single HydrateFunc which does not have a [HydrateConfig]:
-		DefaultConcurrency: &plugin.DefaultConcurrencyConfig{
-			DefaultMaxConcurrency: 100,
-		}
-		
+
+	DefaultConcurrency: &plugin.DefaultConcurrencyConfig{
+		DefaultMaxConcurrency: 100,
+	}
+
 Do both:
-		DefaultConcurrency: &plugin.DefaultConcurrencyConfig{
-			TotalMaxConcurrency:   500,
-			DefaultMaxConcurrency: 200,
-		}
+
+	DefaultConcurrency: &plugin.DefaultConcurrencyConfig{
+		TotalMaxConcurrency:   500,
+		DefaultMaxConcurrency: 200,
+	}
 
 Plugin examples:
   - [hackernews]
@@ -31,14 +34,14 @@ Plugin examples:
 */
 type DefaultConcurrencyConfig struct {
 	// sets how many HydrateFunc calls can run concurrently in total
-	TotalMaxConcurrency   int
+	TotalMaxConcurrency int
 	// sets the default for how many calls to each HydrateFunc can run concurrently
 	DefaultMaxConcurrency int
 }
 
 // concurrencyManager struct ensures that hydrate functions stay within concurrency limits
 type concurrencyManager struct {
-	mut sync.Mutex
+	mut sync.RWMutex
 	// the maximum number of all hydrate calls which can run concurrently
 	maxConcurrency int
 	// the maximum concurrency for a single hydrate call
@@ -80,25 +83,25 @@ func newConcurrencyManager(t *Table) *concurrencyManager {
 // StartIfAllowed checks whether the named hydrate call is permitted to start
 // based on the number of running instances of that call, and the total calls in progress
 func (c *concurrencyManager) StartIfAllowed(name string, maxCallConcurrency int) (res bool) {
-	c.mut.Lock()
-	defer c.mut.Unlock()
+	// acquire a Read lock
+	c.mut.RLock()
+	// how many concurrent executions of this function are in progress right now?
+	currentExecutions := c.callMap[name]
+	// ensure we unlock
+	c.mut.RUnlock()
 
-	// is the total call limit exceeded?
-	if c.maxConcurrency > 0 && c.callsInProgress == c.maxConcurrency {
+	if !c.canStart(currentExecutions, maxCallConcurrency) {
 		return false
 	}
 
-	// if there is no config or empty config, the maxCallConcurrency will be 0
-	// - use defaultMaxConcurrencyPerCall set on the concurrencyManager
-	if maxCallConcurrency == 0 {
-		maxCallConcurrency = c.defaultMaxConcurrencyPerCall
-	}
-
-	// how many concurrent executions of this function are in progress right now?
-	currentExecutions := c.callMap[name]
+	// upgrade the mutex to a Write lock
+	c.mut.Lock()
+	// ensure we unlock
+	defer c.mut.Unlock()
 
-	// if we at the call limit return
-	if maxCallConcurrency > 0 && currentExecutions == maxCallConcurrency {
+	// check again in case another thread grabbed the Write lock before us
+	currentExecutions = c.callMap[name]
+	if !c.canStart(currentExecutions, maxCallConcurrency) {
 		return false
 	}
 
@@ -113,6 +116,26 @@ func (c *concurrencyManager) StartIfAllowed(name string, maxCallConcurrency int)
 	if c.callsInProgress > c.maxCallsInProgress {
 		c.maxCallsInProgress = c.callsInProgress
 	}
+
+	return true
+}
+
+func (c *concurrencyManager) canStart(currentExecutions int, maxCallConcurrency int) bool {
+	// is the total call limit exceeded?
+	if c.maxConcurrency > 0 && c.callsInProgress == c.maxConcurrency {
+		return false
+	}
+
+	// if there is no config or empty config, the maxCallConcurrency will be 0
+	// - use defaultMaxConcurrencyPerCall set on the concurrencyManager
+	if maxCallConcurrency == 0 {
+		maxCallConcurrency = c.defaultMaxConcurrencyPerCall
+	}
+
+	// if we at the call limit return
+	if maxCallConcurrency > 0 && currentExecutions == maxCallConcurrency {
+		return false
+	}
 	return true
 }
 
@@ -123,10 +146,11 @@ func (c *concurrencyManager) Finished(name string) {
 			log.Printf("[WARN] %v", r)
 		}
 	}()
+	// acquire a Write lock
 	c.mut.Lock()
-	defer c.mut.Unlock()
 	c.callMap[name]--
 	c.callsInProgress--
+	c.mut.Unlock()
 }
 
 // Close executes when the query is complete and dumps out the concurrency stats

plugin/concurrency.go

@@ -137,81 +137,6 @@ type extraPropertyWithAnnotation struct {
 
 var testCasesParseConfig = map[string]parseConfigTest{
 	// cty tag implementation
-	"struct property cty": {
-		source: `
-	name = { name = "foo"}
-	`,
-		connectionConfigSchema: &ConnectionConfigSchema{
-			NewInstance: func() interface{} { return &structPropertyCty{} },
-			Schema: map[string]*schema.Attribute{
-				"name": {
-					Type: schema.TypeMap,
-					AttrTypes: map[string]*schema.Attribute{
-						"name": {
-							Type: schema.TypeString,
-						},
-					},
-				},
-			},
-		},
-		expected: structPropertyCty{
-			childStructCty{"foo"}},
-	},
-	"struct slice property cty": {
-		source: `
-  tables = [
-    {
-      name    = "test1"
-      columns = [
-        {
-          name = "c1"
-          type = "string"
-        },
-        {
-          name = "c2"
-          type = "string"
-        },
-      ]
-    }
-  ]
-	`,
-		connectionConfigSchema: &ConnectionConfigSchema{
-			NewInstance: func() interface{} { return &structSlicePropertyCty{} },
-			Schema: map[string]*schema.Attribute{
-				"tables": {
-					Type: schema.TypeList,
-					Elem: &schema.Attribute{
-						Type: schema.TypeMap,
-						AttrTypes: map[string]*schema.Attribute{
-							"name": {
-								Type: schema.TypeString,
-							},
-							"columns": {
-								Type: schema.TypeList,
-								Elem: &schema.Attribute{
-
-									Type: schema.TypeMap,
-									AttrTypes: map[string]*schema.Attribute{
-										"name": {
-											Type: schema.TypeString,
-										},
-										"type": {
-											Type: schema.TypeString,
-										},
-									},
-								},
-							},
-						},
-					},
-				},
-			},
-		},
-		expected: structSlicePropertyCty{
-			Tables: []configTableCty{{
-				Name:    "test1",
-				Columns: []configColumnCty{{"c1", "string"}, {"c2", "string"}},
-			}}},
-	},
 	"array property cty": {
 		source: `
 	regions = ["us-east-1","us-west-2"]