Define the level of concurrency from configuration file.

Author: clems4ever

cmd/go-graphkb/config.yml

@@ -6,3 +6,6 @@ mariadb_username: graphkb
 mariadb_password: password
 mariadb_host: db
 mariadb_database: graphkb
+
+# The level of concurrency allowed by the graph update API.
+concurrency: 32

cmd/go-graphkb/main.go

@@ -118,8 +118,12 @@ func listen(cmd *cobra.Command, args []string) {
 	}
 
 	listenInterface := viper.GetString("server_listen")
+	concurrency := viper.GetInt64("concurrency")
+	if concurrency == 0 {
+		concurrency = 32
+	}
 
-	server.StartServer(listenInterface, Database, Database, Database, Database)
+	server.StartServer(listenInterface, Database, Database, Database, Database, concurrency)
 }
 
 func read(cmd *cobra.Command, args []string) {

internal/client/transaction.go

@@ -50,12 +50,10 @@ func withRetryOnTooManyRequests(fn func() error, maxRetries int) error {
 	trials := 0
 	for {
 		err := fn()
-		if err == ErrTooManyRequests {
+		if err != nil {
 			backoffTime := time.Duration(int(math.Pow(1.01, float64(trials)))*15) * time.Second
 			fmt.Printf("Sleeping for %d seconds", backoffTime/time.Second)
 			time.Sleep(backoffTime)
-		} else if err != nil {
-			return err
 		} else {
 			return nil
 		}
@@ -66,38 +64,6 @@ func withRetryOnTooManyRequests(fn func() error, maxRetries int) error {
 	}
 }
 
-func runTasksInParallel(fn func(data interface{}) error, workers int, items []interface{}) error {
-	var wg sync.WaitGroup
-	var err error
-
-	dataC := make(chan interface{})
-
-	wg.Add(workers)
-
-	for i := 0; i < workers; i++ {
-		go func() {
-			for {
-				select {
-				case d, more := <-dataC:
-					// if d is nil, it means the channel has been closed
-					if !more || err != nil {
-						return
-					}
-					err = fn(d)
-				}
-			}
-		}()
-	}
-
-	for _, d := range items {
-		dataC <- d
-	}
-
-	close(dataC)
-	wg.Wait()
-	return err
-}
-
 // Commit commit the transaction and gives ownership to the source for caching.
 func (cgt *Transaction) Commit() (*knowledge.Graph, error) {
 	sg := cgt.newGraph.ExtractSchema()

internal/handlers/handler_update_graph.go

@@ -12,9 +12,7 @@ import (
 	"golang.org/x/sync/semaphore"
 )
 
-var updatesSemaphore = semaphore.NewWeighted(int64(256))
-
-func handleUpdate(registry sources.Registry, fn func(source string, body io.Reader) error) http.HandlerFunc {
+func handleUpdate(registry sources.Registry, fn func(source string, body io.Reader) error, sem *semaphore.Weighted) http.HandlerFunc {
 	return func(w http.ResponseWriter, r *http.Request) {
 		ok, source, err := IsTokenValid(registry, r)
 		if err != nil {
@@ -28,12 +26,12 @@ func handleUpdate(registry sources.Registry, fn func(source string, body io.Read
 		}
 
 		{
-			ok = updatesSemaphore.TryAcquire(1)
+			ok = sem.TryAcquire(1)
 			if !ok {
 				ReplyWithTooManyRequests(w)
 				return
 			}
-			defer updatesSemaphore.Release(1)
+			defer sem.Release(1)
 
 			if err = fn(source, r.Body); err != nil {
 				ReplyWithInternalError(w, err)
@@ -50,7 +48,7 @@ func handleUpdate(registry sources.Registry, fn func(source string, body io.Read
 }
 
 // PutSchema upsert an asset into the graph of the data source
-func PutSchema(registry sources.Registry, graphUpdater *knowledge.GraphUpdater) http.HandlerFunc {
+func PutSchema(registry sources.Registry, graphUpdater *knowledge.GraphUpdater, sem *semaphore.Weighted) http.HandlerFunc {
 	return handleUpdate(registry, func(source string, body io.Reader) error {
 		requestBody := client.PutGraphSchemaRequestBody{}
 		if err := json.NewDecoder(body).Decode(&requestBody); err != nil {
@@ -60,11 +58,11 @@ func PutSchema(registry sources.Registry, graphUpdater *knowledge.GraphUpdater)
 		// TODO(c.michaud): verify compatibility of the schema with graph updates
 		graphUpdater.UpdateSchema(source, requestBody.Schema)
 		return nil
-	})
+	}, sem)
 }
 
 // PutAsset upsert an asset into the graph of the data source
-func PutAsset(registry sources.Registry, graphUpdater *knowledge.GraphUpdater) http.HandlerFunc {
+func PutAsset(registry sources.Registry, graphUpdater *knowledge.GraphUpdater, sem *semaphore.Weighted) http.HandlerFunc {
 	return handleUpdate(registry, func(source string, body io.Reader) error {
 		requestBody := client.PutGraphAssetRequestBody{}
 		if err := json.NewDecoder(body).Decode(&requestBody); err != nil {
@@ -74,11 +72,11 @@ func PutAsset(registry sources.Registry, graphUpdater *knowledge.GraphUpdater) h
 		// TODO(c.michaud): verify compatibility of the schema with graph updates
 		graphUpdater.UpsertAsset(source, requestBody.Asset)
 		return nil
-	})
+	}, sem)
 }
 
 // PutRelation upsert a relation into the graph of the data source
-func PutRelation(registry sources.Registry, graphUpdater *knowledge.GraphUpdater) http.HandlerFunc {
+func PutRelation(registry sources.Registry, graphUpdater *knowledge.GraphUpdater, sem *semaphore.Weighted) http.HandlerFunc {
 	return handleUpdate(registry, func(source string, body io.Reader) error {
 		requestBody := client.PutGraphRelationRequestBody{}
 		if err := json.NewDecoder(body).Decode(&requestBody); err != nil {
@@ -88,11 +86,11 @@ func PutRelation(registry sources.Registry, graphUpdater *knowledge.GraphUpdater
 		// TODO(c.michaud): verify compatibility of the schema with graph updates
 		graphUpdater.UpsertRelation(source, requestBody.Relation)
 		return nil
-	})
+	}, sem)
 }
 
 // DeleteAsset delete an asset from the graph of the data source
-func DeleteAsset(registry sources.Registry, graphUpdater *knowledge.GraphUpdater) http.HandlerFunc {
+func DeleteAsset(registry sources.Registry, graphUpdater *knowledge.GraphUpdater, sem *semaphore.Weighted) http.HandlerFunc {
 	return handleUpdate(registry, func(source string, body io.Reader) error {
 		requestBody := client.DeleteGraphAssetRequestBody{}
 		if err := json.NewDecoder(body).Decode(&requestBody); err != nil {
@@ -102,11 +100,11 @@ func DeleteAsset(registry sources.Registry, graphUpdater *knowledge.GraphUpdater
 		// TODO(c.michaud): verify compatibility of the schema with graph updates
 		graphUpdater.RemoveAsset(source, requestBody.Asset)
 		return nil
-	})
+	}, sem)
 }
 
 // DeleteRelation upsert a relation into the graph of the data source
-func DeleteRelation(registry sources.Registry, graphUpdater *knowledge.GraphUpdater) http.HandlerFunc {
+func DeleteRelation(registry sources.Registry, graphUpdater *knowledge.GraphUpdater, sem *semaphore.Weighted) http.HandlerFunc {
 	return handleUpdate(registry, func(source string, body io.Reader) error {
 		requestBody := client.DeleteGraphRelationRequestBody{}
 		if err := json.NewDecoder(body).Decode(&requestBody); err != nil {
@@ -116,5 +114,5 @@ func DeleteRelation(registry sources.Registry, graphUpdater *knowledge.GraphUpda
 		// TODO(c.michaud): verify compatibility of the schema with graph updates
 		graphUpdater.RemoveRelation(source, requestBody.Relation)
 		return nil
-	})
+	}, sem)
 }

internal/server/server.go

@@ -17,6 +17,7 @@ import (
 	"github.com/clems4ever/go-graphkb/internal/sources"
 	"github.com/clems4ever/go-graphkb/internal/utils"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
+	"golang.org/x/sync/semaphore"
 
 	auth "github.com/abbot/go-http-auth"
 
@@ -276,7 +277,8 @@ func StartServer(listenInterface string,
 	database knowledge.GraphDB,
 	schemaPersistor schema.Persistor,
 	sourcesRegistry sources.Registry,
-	queryHistorizer history.Historizer) {
+	queryHistorizer history.Historizer,
+	concurrency int64) {
 
 	r := mux.NewRouter()
 
@@ -314,11 +316,13 @@ func StartServer(listenInterface string,
 
 	r.HandleFunc("/api/graph/read", getGraphRead(sourcesRegistry, database)).Methods("GET")
 
-	r.HandleFunc("/api/graph/schema", handlers.PutSchema(sourcesRegistry, graphUpdater)).Methods("PUT")
-	r.HandleFunc("/api/graph/asset", handlers.PutAsset(sourcesRegistry, graphUpdater)).Methods("PUT")
-	r.HandleFunc("/api/graph/relation", handlers.PutRelation(sourcesRegistry, graphUpdater)).Methods("PUT")
-	r.HandleFunc("/api/graph/asset", handlers.DeleteAsset(sourcesRegistry, graphUpdater)).Methods("DELETE")
-	r.HandleFunc("/api/graph/relation", handlers.DeleteRelation(sourcesRegistry, graphUpdater)).Methods("DELETE")
+	sem := semaphore.NewWeighted(concurrency)
+
+	r.HandleFunc("/api/graph/schema", handlers.PutSchema(sourcesRegistry, graphUpdater, sem)).Methods("PUT")
+	r.HandleFunc("/api/graph/asset", handlers.PutAsset(sourcesRegistry, graphUpdater, sem)).Methods("PUT")
+	r.HandleFunc("/api/graph/relation", handlers.PutRelation(sourcesRegistry, graphUpdater, sem)).Methods("PUT")
+	r.HandleFunc("/api/graph/asset", handlers.DeleteAsset(sourcesRegistry, graphUpdater, sem)).Methods("DELETE")
+	r.HandleFunc("/api/graph/relation", handlers.DeleteRelation(sourcesRegistry, graphUpdater, sem)).Methods("DELETE")
 
 	r.HandleFunc("/api/query", postQueryHandler).Methods("POST")
 	r.PathPrefix("/").Handler(http.FileServer(http.Dir("./web/build/")))
@@ -327,12 +331,12 @@ func StartServer(listenInterface string,
 
 	var err error
 	if viper.GetString("server_tls_cert") != "" {
-		fmt.Printf("Listening on %s with TLS enabled, the connection is secure\n", listenInterface)
+		fmt.Printf("Listening on %s with TLS enabled, the connection is secure [concurrency=%d]\n", listenInterface, concurrency)
 		err = http.ListenAndServeTLS(listenInterface, viper.GetString("server_tls_cert"),
 			viper.GetString("server_tls_key"), r)
 	} else {
-		fmt.Printf("[WARNING] Listening on %s with TLS disabled. Use `server_tls_cert` option to setup a certificate\n",
-			listenInterface)
+		fmt.Printf("[WARNING] Listening on %s with TLS disabled. Use `server_tls_cert` option to setup a certificate [concurrency=%d]\n",
+			listenInterface, concurrency)
 		err = http.ListenAndServe(listenInterface, r)
 	}
 	if err != nil {