Athens: introduce Single Flight

cmd/olympus/actions/app.go

@@ -15,12 +15,14 @@ import (
 	"github.com/gomods/athens/pkg/cdn/metadata/azurecdn"
 	"github.com/gomods/athens/pkg/config/env"
 	"github.com/gomods/athens/pkg/download"
-	"github.com/gomods/athens/pkg/download/goget"
 	"github.com/gomods/athens/pkg/eventlog"
 	"github.com/gomods/athens/pkg/log"
+	"github.com/gomods/athens/pkg/module"
+	"github.com/gomods/athens/pkg/stash"
 	"github.com/gomods/athens/pkg/storage"
 	"github.com/gomodule/redigo/redis"
 	"github.com/rs/cors"
+	"github.com/spf13/afero"
 	"github.com/unrolled/secure"
 )
 
@@ -85,6 +87,7 @@ func App(config *AppConfig) (*buffalo.App, error) {
 		WorkerOff:   true, // TODO(marwan): turned off until worker is being used.
 		Logger:      blggr,
 	})
+
 	// Automatically redirect to SSL
 	app.Use(ssl.ForceSSL(secure.Options{
 		SSLRedirect:     ENV == "production",
@@ -123,13 +126,23 @@ func App(config *AppConfig) (*buffalo.App, error) {
 	app.GET("/healthz", healthHandler)
 
 	// Download Protocol
-	gg, err := goget.New()
+	goBin := env.GoBinPath()
+	fs := afero.NewOsFs()
+	mf, err := module.NewGoGetFetcher(goBin, fs)
 	if err != nil {
 		return nil, err
 	}
-	dp := download.New(gg, config.Storage, env.GoGetWorkers())
-	opts := &download.HandlerOpts{Protocol: dp, Logger: lggr, Engine: renderEng}
-	download.RegisterHandlers(app, opts)
+	st := stash.New(mf, config.Storage)
+	dpOpts := &download.Opts{
+		Storage:   config.Storage,
+		Stasher:   st,
+		GoBinPath: goBin,
+		Fs:        fs,
+	}
+	dp := download.New(dpOpts)
+
+	handlerOpts := &download.HandlerOpts{Protocol: dp, Logger: lggr, Engine: renderEng}
+	download.RegisterHandlers(app, handlerOpts)
 
 	app.ServeFiles("/", assetsBox) // serve files from the public directory
 

cmd/proxy/actions/app_proxy.go

@@ -4,9 +4,12 @@ import (
 	"github.com/gobuffalo/buffalo"
 	"github.com/gomods/athens/pkg/config/env"
 	"github.com/gomods/athens/pkg/download"
-	"github.com/gomods/athens/pkg/download/goget"
+	"github.com/gomods/athens/pkg/download/addons"
 	"github.com/gomods/athens/pkg/log"
+	"github.com/gomods/athens/pkg/module"
+	"github.com/gomods/athens/pkg/stash"
 	"github.com/gomods/athens/pkg/storage"
+	"github.com/spf13/afero"
 )
 
 func addProxyRoutes(
@@ -18,13 +21,45 @@ func addProxyRoutes(
 	app.GET("/healthz", healthHandler)
 
 	// Download Protocol
-	gg, err := goget.New()
+	// the download.Protocol and the stash.Stasher interfaces are composable
+	// in a middleware fashion. Therefore you can separate concerns
+	// by the functionality: a download.Protocol that just takes care
+	// of "go getting" things, and another Protocol that just takes care
+	// of "pooling" requests etc.
+
+	// In our case, we'd like to compose both interfaces in a particular
+	// order to ensure logical ordering of execution.
+
+	// Here's the order of an incoming request to the download.Protocol:
+
+	// 1. The downloadpool gets hit first, and manages concurrent requests
+	// 2. The downloadpool passes the request to its parent Protocol: stasher
+	// 3. The stasher Protocol checks storage first, and if storage is empty
+	// it makes a Stash request to the stash.Stasher interface.
+
+	// Once the stasher picks up an order, here's how the requests go in order:
+	// 1. The singleflight picks up the first request and latches duplicate ones.
+	// 2. The singleflight passes the stash to its parent: stashpool.
+	// 3. The stashpool manages limiting concurrent requests and passes them to stash.
+	// 4. The plain stash.New just takes a request from upstream and saves it into storage.
+	goBin := env.GoBinPath()
+	fs := afero.NewOsFs()
+	mf, err := module.NewGoGetFetcher(goBin, fs)
 	if err != nil {
 		return err
 	}
-	p := download.New(gg, s, env.GoGetWorkers())
-	opts := &download.HandlerOpts{Protocol: p, Logger: l, Engine: proxy}
-	download.RegisterHandlers(app, opts)
+	st := stash.New(mf, s, stash.WithPool(env.GoGetWorkers()), stash.WithSingleflight)
+
+	dpOpts := &download.Opts{
+		Storage:   s,
+		Stasher:   st,
+		GoBinPath: goBin,
+		Fs:        fs,
+	}
+	dp := download.New(dpOpts, addons.WithPool(env.ProtocolWorkers()))
+
+	handlerOpts := &download.HandlerOpts{Protocol: dp, Logger: l, Engine: proxy}
+	download.RegisterHandlers(app, handlerOpts)
 
 	return nil
 }

pkg/config/env/go.go

@@ -36,3 +36,24 @@ func GoGetWorkers() int {
 
 	return num
 }
+
+// ProtocolWorkers returns how many concurrent
+// requests can you handle at a time for all
+// download protocol paths. This is different from
+// GoGetWorkers in that you can potentially serve
+// 30 requests to the Download Protocol but only 5
+// at a time can stash a module from Upstream to Storage.
+func ProtocolWorkers() int {
+	defaultNum := 30
+	str := os.Getenv("ATHENS_PROTOCOL_WORKERS")
+	if str == "" {
+		return defaultNum
+	}
+
+	num, err := strconv.Atoi(str)
+	if err != nil {
+		return defaultNum
+	}
+
+	return num
+}

pkg/download/addons/with_pool.go

@@ -0,0 +1,129 @@
+package addons
+
+import (
+	"context"
+	"io"
+
+	"github.com/gomods/athens/pkg/download"
+	"github.com/gomods/athens/pkg/errors"
+	"github.com/gomods/athens/pkg/storage"
+)
+
+type withpool struct {
+	dp download.Protocol
+
+	// jobCh is a channel that takes an anonymous
+	// function that it executes based on the pool's
+	// business. The design levarages closures
+	// so that the worker does not need to worry about
+	// what the type of job it is taking (Info, Zip etc),
+	// it just regulates functions and executes them
+	// in a worker-pool fashion.
+	jobCh chan func()
+}
+
+// WithPool takes a download Protocol and a number of workers
+// and creates a N worker pool that share all the download.Protocol
+// methods.
+func WithPool(workers int) download.Wrapper {
+	return func(dp download.Protocol) download.Protocol {
+		jobCh := make(chan func())
+		p := &withpool{dp: dp, jobCh: jobCh}
+
+		p.start(workers)
+		return p
+	}
+}
+
+func (p *withpool) start(numWorkers int) {
+	for i := 0; i < numWorkers; i++ {
+		go p.listen()
+	}
+}
+
+func (p *withpool) listen() {
+	for f := range p.jobCh {
+		f()
+	}
+}
+
+func (p *withpool) List(ctx context.Context, mod string) ([]string, error) {
+	const op errors.Op = "pool.List"
+	var vers []string
+	var err error
+	done := make(chan struct{}, 1)
+	p.jobCh <- func() {
+		vers, err = p.dp.List(ctx, mod)
+		close(done)
+	}
+	<-done
+	if err != nil {
+		return nil, errors.E(op, err)
+	}
+
+	return vers, nil
+}
+
+func (p *withpool) Info(ctx context.Context, mod, ver string) ([]byte, error) {
+	const op errors.Op = "pool.Info"
+	var info []byte
+	var err error
+	done := make(chan struct{}, 1)
+	p.jobCh <- func() {
+		info, err = p.dp.Info(ctx, mod, ver)
+		close(done)
+	}
+	<-done
+	if err != nil {
+		return nil, errors.E(op, err)
+	}
+	return info, nil
+}
+
+func (p *withpool) Latest(ctx context.Context, mod string) (*storage.RevInfo, error) {
+	const op errors.Op = "pool.Latest"
+	var info *storage.RevInfo
+	var err error
+	done := make(chan struct{}, 1)
+	p.jobCh <- func() {
+		info, err = p.dp.Latest(ctx, mod)
+		close(done)
+	}
+	<-done
+	if err != nil {
+		return nil, errors.E(op, err)
+	}
+	return info, nil
+}
+
+func (p *withpool) GoMod(ctx context.Context, mod, ver string) ([]byte, error) {
+	const op errors.Op = "pool.GoMod"
+	var goMod []byte
+	var err error
+	done := make(chan struct{}, 1)
+	p.jobCh <- func() {
+		goMod, err = p.dp.GoMod(ctx, mod, ver)
+		close(done)
+	}
+	<-done
+	if err != nil {
+		return nil, errors.E(op, err)
+	}
+	return goMod, nil
+}
+
+func (p *withpool) Zip(ctx context.Context, mod, ver string) (io.ReadCloser, error) {
+	const op errors.Op = "pool.Zip"
+	var zip io.ReadCloser
+	var err error
+	done := make(chan struct{}, 1)
+	p.jobCh <- func() {
+		zip, err = p.dp.Zip(ctx, mod, ver)
+		close(done)
+	}
+	<-done
+	if err != nil {
+		return nil, errors.E(op, err)
+	}
+	return zip, nil
+}