Add mechanism to limit max hydrate concurrency. Closes #12

Author: tyagiparth

plugin/concurrency.go

@@ -0,0 +1,93 @@
+package plugin
+
+import (
+	"log"
+	"sync"
+)
+
+// if no max concurrency is specified in the plugin, use this value
+const defaultMaxConcurrency = 1000
+
+// if no max call concurrency is specified for a hydrate function, use this value
+const defaultMaxConcurrencyPerCall = 500
+
+// ConcurrencyManager :: struct which ensures hydrate funcitons stay within concurrency limits
+type ConcurrencyManager struct {
+	mut sync.Mutex
+	// the maximun number of all hydrate calls which can run concurrently
+	maxConcurrency int
+	// the maximum concurrency for a single hydrate call
+	// (this may be overridden by the HydrateConfig for the call)
+	defaultMaxConcurrencyPerCall int
+	// total number of hydrate calls in progress
+	callsInProgress int
+	// map of the number of instances of each call in progress
+	callMap map[string]int
+}
+
+func newConcurrencyManager(t *Table) *ConcurrencyManager {
+	// if plugin does not define max concurrency, use default
+	max := defaultMaxConcurrency
+	// if hydrate calls do not define max concurrency, use default
+	maxPerCall := defaultMaxConcurrencyPerCall
+	if config := t.Plugin.DefaultHydrateConfig; config != nil {
+		if config.MaxConcurrency != 0 {
+			max = config.MaxConcurrency
+		}
+		if config.DefaultMaxConcurrencyPerCall != 0 {
+			maxPerCall = config.DefaultMaxConcurrencyPerCall
+		} else if max < maxPerCall {
+			// if the default call concurrency is greater than the toal max concurrency, clamp to total
+			maxPerCall = max
+		}
+	}
+	return &ConcurrencyManager{
+		maxConcurrency:               max,
+		defaultMaxConcurrencyPerCall: maxPerCall,
+		callMap:                      make(map[string]int),
+	}
+}
+
+// StartIfAllowed :: check 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()
+
+	// is the total call limit exceeded?
+	if 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
+	}
+
+	// how many concurrent executions of this function ar in progress right now?
+	currentExecutions := c.callMap[name]
+
+	// if we at the call limit return
+	if currentExecutions == maxCallConcurrency {
+		return false
+	}
+
+	// to get here we are allowed to execute - increment the call counters
+	c.callMap[name] = currentExecutions + 1
+	c.callsInProgress++
+	return true
+}
+
+// Finished :: decrement the counter for the named function
+func (c *ConcurrencyManager) Finished(name string) {
+	defer func() {
+		if r := recover(); r != nil {
+			log.Printf("[WARN] %v", r)
+		}
+	}()
+	c.mut.Lock()
+	defer c.mut.Unlock()
+	c.callMap[name]--
+	c.callsInProgress--
+}

plugin/hydrate.go

@@ -6,50 +6,83 @@ import (
 	"github.com/turbot/go-kit/helpers"
 )
 
+// HydrateData :: the input data passed to every hydrate function
 type HydrateData struct {
 	Item           interface{}
-	Params         map[string]string
 	HydrateResults map[string]interface{}
 }
 
-// perform shallow clone
-func (h *HydrateData) Clone() *HydrateData {
-	return &HydrateData{
-		Item:           h.Item,
-		Params:         h.Params,
-		HydrateResults: h.HydrateResults,
-	}
-}
-
-// HydrateFunc is a function which retrieves some or all row data for a single row item.
+// HydrateFunc :: a function which retrieves some or all row data for a single row item.
 type HydrateFunc func(context.Context, *QueryData, *HydrateData) (interface{}, error)
 
 // HydrateDependencies :: define the hydrate function dependencies - other hydrate functions which must be run first
+// Deprecated: used HydrateConfig
 type HydrateDependencies struct {
 	Func    HydrateFunc
 	Depends []HydrateFunc
 }
 
+// HydrateConfig :: define the hydrate function configurations, Name, Maximum number of concurrent calls to be allowed, dependencies
+type HydrateConfig struct {
+	Func           HydrateFunc
+	MaxConcurrency int
+	// ConcurrencyMapKey ConcurrencyMapKeyFunc
+	// ShouldRetryError ErrorPredicate
+	// ShouldIgnoreError ErrorPredicate
+	Depends []HydrateFunc
+}
+
+// DefaultHydrateConfig :: plugin levelk config to define default hydrate concurrency
+// - used if no HydrateConfig is specified for a specific call
+type DefaultHydrateConfig struct {
+	// max number of ALL hydrate calls in progress
+	MaxConcurrency               int
+	DefaultMaxConcurrencyPerCall int
+}
+
+// HydrateCall :: struct encapsulating a hydrate call, its config and dependencies
 type HydrateCall struct {
 	Func HydrateFunc
 	// the dependencies expressed using function name
 	Depends []string
+	Config  *HydrateConfig
 }
 
-func newHydrateCall(hydrateFunc HydrateFunc, dependencies []HydrateFunc) *HydrateCall {
-	res := &HydrateCall{Func: hydrateFunc}
-	for _, f := range dependencies {
+func newHydrateCall(hydrateFunc HydrateFunc, config *HydrateConfig) *HydrateCall {
+	res := &HydrateCall{Func: hydrateFunc, Config: config}
+	for _, f := range config.Depends {
 		res.Depends = append(res.Depends, helpers.GetFunctionName(f))
 	}
 	return res
 }
 
-// CanStart :: can this hydrate call - check whether all dependency hydrate functions have been completed
-func (h HydrateCall) CanStart(rowData *RowData) bool {
+// CanStart :: return whether this hydrate call can execute
+// - check whether all dependency hydrate functions have been completed
+// - check whether the concurrency limits would be exceeded
+
+func (h HydrateCall) CanStart(rowData *RowData, name string, concurrencyManager *ConcurrencyManager) bool {
 	for _, dep := range h.Depends {
 		if !helpers.StringSliceContains(rowData.getHydrateKeys(), dep) {
 			return false
 		}
 	}
-	return true
+	// ask the concurrency manager whether the call can start
+	// NOTE: if the call is allowed to start, the concurrency manager ASSUMES THE CALL WILL START
+	// and increments the counters
+	// it may seem more logical to do this in the Start() function below, but we need to check and increment the counters
+	// within the same mutex lock to ensure another call does not start between checking and starting
+	return concurrencyManager.StartIfAllowed(name, h.Config.MaxConcurrency)
+}
+
+// Start :: start a hydrate call
+func (h *HydrateCall) Start(ctx context.Context, r *RowData, hydrateFuncName string, concurrencyManager *ConcurrencyManager) {
+	// tell the roewdata to wait for this call to complete
+	r.wg.Add(1)
+
+	// call callHydrate async, ignoring return values
+	go func() {
+		r.callHydrate(ctx, r.queryData, h.Func, hydrateFuncName)
+		// decrement number of hydrate functions running
+		concurrencyManager.Finished(hydrateFuncName)
+	}()
 }

plugin/plugin.go

@@ -21,6 +21,7 @@ type Plugin struct {
 	TableMap         map[string]*Table
 	DefaultTransform *transform.ColumnTransforms
 	DefaultGetConfig *GetConfig
+	DefaultHydrateConfig *DefaultHydrateConfig
 	// every table must implement these columns
 	RequiredColumns []*Column
 }
@@ -96,6 +97,7 @@ func (p *Plugin) Execute(req *proto.ExecuteRequest, stream proto.WrapperPlugin_E
 	ctx := context.WithValue(context.Background(), ContextKeyLogger, p.Logger)
 	log.Printf("[TRACE] calling fetchItems, table: %s\n", table.Name)
 
+	// asyncronously fetch items
 	table.fetchItems(ctx, d)
 
 	log.Println("[TRACE] after fetchItems")

plugin/query_data.go

@@ -2,6 +2,7 @@ package plugin
 
 import (
 	"context"
+	"fmt"
 	"log"
 
 	"github.com/turbot/go-kit/helpers"
@@ -29,10 +30,11 @@ type QueryData struct {
 	QueryContext *pb.QueryContext
 
 	// internal
-	hydrateCalls []*HydrateCall
-	rowDataChan  chan *RowData
-	errorChan    chan error
-	stream       pb.WrapperPlugin_ExecuteServer
+	hydrateCalls       []*HydrateCall
+	concurrencyManager *ConcurrencyManager
+	rowDataChan        chan *RowData
+	errorChan          chan error
+	stream             pb.WrapperPlugin_ExecuteServer
 	// wait group used to syncronise parent-child list fetches - each child hydrate function increments this wait group
 	listWg sync.WaitGroup
 }
@@ -56,6 +58,7 @@ func newQueryData(queryContext *pb.QueryContext, table *Table, stream pb.Wrapper
 	ensureColumns(queryContext, table)
 
 	d.hydrateCalls = table.requiredHydrateCalls(queryContext.Columns, d.FetchType)
+	d.concurrencyManager = newConcurrencyManager(table)
 	return d
 }
 
@@ -232,6 +235,9 @@ func (d *QueryData) buildRows(ctx context.Context) chan *pb.Row {
 // execute necessary hydrate calls to populate row data
 func (d *QueryData) buildRow(ctx context.Context, rowData *RowData, rowChan chan *pb.Row, wg *sync.WaitGroup) {
 	defer func() {
+		if r := recover(); r != nil {
+			d.streamError(ToError(r))
+		}
 		wg.Done()
 	}()
 
@@ -268,3 +274,13 @@ func (d *QueryData) singleEqualsQual(column string) (*pb.Qual, bool) {
 	}
 	return nil, false
 }
+
+// remove once go-kit version 0.2.0 is released
+// ToError :: if supplied value is already an error, return it, otherwise format it as an error
+func ToError(val interface{}) error {
+	if e, ok := val.(error); ok {
+		return e
+	} else {
+		return fmt.Errorf("%v", val)
+	}
+}

plugin/required_hydrate_calls.go

@@ -1,8 +1,9 @@
 package plugin
 
 import (
-	"github.com/turbot/go-kit/helpers"
 	"log"
+
+	"github.com/turbot/go-kit/helpers"
 )
 
 // helper class to build list of required hydrate calls
@@ -30,12 +31,18 @@ func (c requiredHydrateCallBuilder) Add(hydrateFunc HydrateFunc) {
 
 		log.Printf("[TRACE] adding hydration function '%s' to hydrationMap\n", hydrateName)
 
-		// get any dependencies for this hydrate function
-		dependencies := c.table.getHydrateDependencies(hydrateName)
-		c.requiredHydrateCalls[hydrateName] = newHydrateCall(hydrateFunc, dependencies)
+		// get the config for this hydrate function
+		config := c.table.getHydrateConfig(hydrateName)
+
+		// get any dependencies for this hydrate function. if no hydrate dependencies are specified in the hydrate config, check the deprecated "HydrateDependencies" property
+		if config.Depends == nil {
+			config.Depends = c.table.getHydrateDependencies(hydrateName)
+		}
+
+		c.requiredHydrateCalls[hydrateName] = newHydrateCall(hydrateFunc, config)
 
 		// now add dependencies (we have already checked for circular dependencies so recursion is fine
-		for _, dep := range dependencies {
+		for _, dep := range config.Depends {
 			c.Add(dep)
 		}
 	}

plugin/row.go

@@ -52,21 +52,18 @@ func (r *RowData) getRow(ctx context.Context) (*pb.Row, error) {
 	// - these populate the row with data entries corresponding to the hydrate function nameSP_LOG=TRACE
 	// keep looping round hydrate functions until they are all started
 
-	// make a map of started hydrate calls
+	// make a map of started hydrate calls for this row - this is used the determine which calls have not started yet
 	var callsStarted = map[string]bool{}
 
 	for {
 		var allStarted = true
 		for _, call := range r.queryData.hydrateCalls {
 			hydrateFuncName := helpers.GetFunctionName(call.Func)
 			if !callsStarted[hydrateFuncName] {
-
-				if call.CanStart(r) {
-					r.wg.Add(1)
-					// call callHydrate async, ignoring return values
-					go r.callHydrate(ctx, r.queryData, call.Func, hydrateFuncName)
+				if call.CanStart(r, hydrateFuncName, r.queryData.concurrencyManager) {
+					// execute the hydrate call asynchronously
+					call.Start(ctx, r, hydrateFuncName, r.queryData.concurrencyManager)
 					callsStarted[hydrateFuncName] = true
-
 				} else {
 					allStarted = false
 				}
@@ -159,6 +156,7 @@ func (r *RowData) callHydrate(ctx context.Context, d *QueryData, hydrateFunc Hyd
 	}
 
 	logging.LogTime(hydrateKey + " end")
+
 	// NOTE: also return the error - is this is being called by as 'get' call we can act on the error immediately
 	return hydrateData, err
 }

plugin/table.go

@@ -22,6 +22,7 @@ type Table struct {
 	Plugin *Plugin
 	// definitions of dependencies between hydrate functions
 	HydrateDependencies []HydrateDependencies
+	HydrateConfig       []HydrateConfig
 }
 
 type GetConfig struct {
@@ -91,3 +92,14 @@ func (t *Table) getHydrateDependencies(hydrateFuncName string) []HydrateFunc {
 	}
 	return []HydrateFunc{}
 }
+
+func (t *Table) getHydrateConfig(hydrateFuncName string) *HydrateConfig {
+	// if a hydrate config is defined see whether this call exists in it
+	for _, d := range t.HydrateConfig {
+		if helpers.GetFunctionName(d.Func) == hydrateFuncName {
+			return &d
+		}
+	}
+	// fallback to return an empty hydrate config
+	return &HydrateConfig{}
+}