Use NATS for shortcut reports in the service.

app/multitenant/common.go

@@ -0,0 +1,22 @@
+package multitenant
+
+import (
+	"time"
+
+	"github.com/prometheus/client_golang/prometheus"
+)
+
+func errorCode(err error) string {
+	if err == nil {
+		return "200"
+	}
+	return "500"
+}
+
+func timeRequest(method string, metric *prometheus.SummaryVec, f func() error) error {
+	startTime := time.Now()
+	err := f()
+	duration := time.Now().Sub(startTime)
+	metric.WithLabelValues(method, errorCode(err)).Observe(float64(duration.Nanoseconds()))
+	return err
+}

app/multitenant/dynamo_collector.go

@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"io"
 	"strconv"
+	"sync"
 	"time"
 
 	log "github.com/Sirupsen/logrus"
@@ -15,6 +16,7 @@ import (
 	"github.com/aws/aws-sdk-go/service/dynamodb"
 	"github.com/aws/aws-sdk-go/service/s3"
 	"github.com/bluele/gcache"
+	"github.com/nats-io/nats"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/ugorji/go/codec"
 	"golang.org/x/net/context"
@@ -24,11 +26,12 @@ import (
 )
 
 const (
-	hourField       = "hour"
-	tsField         = "ts"
-	reportField     = "report"
-	cacheSize       = (15 / 3) * 10 * 5 // (window size * report rate) * number of hosts per user * number of users
-	cacheExpiration = 15 * time.Second
+	hourField             = "hour"
+	tsField               = "ts"
+	reportField           = "report"
+	reportCacheSize       = (15 / 3) * 10 * 5 // (window size * report rate) * number of hosts per user * number of users
+	reportCacheExpiration = 15 * time.Second
+	natsTimeout           = 10 * time.Second
 )
 
 var (
@@ -69,6 +72,12 @@ var (
 		Name:      "s3_request_duration_nanoseconds",
 		Help:      "Time spent doing S3 requests.",
 	}, []string{"method", "status_code"})
+
+	natsRequests = prometheus.NewCounterVec(prometheus.CounterOpts{
+		Namespace: "scope",
+		Name:      "nats_requests",
+		Help:      "Number of NATS requests.",
+	}, []string{"method", "status_code"})
 )
 
 func init() {
@@ -79,6 +88,7 @@ func init() {
 	prometheus.MustRegister(dynamoValueSize)
 	prometheus.MustRegister(reportSize)
 	prometheus.MustRegister(s3RequestDuration)
+	prometheus.MustRegister(natsRequests)
 }
 
 // DynamoDBCollector is a Collector which can also CreateTables
@@ -95,20 +105,53 @@ type dynamoDBCollector struct {
 	bucketName string
 	merger     app.Merger
 	cache      gcache.Cache
+
+	nats        *nats.Conn
+	waitersLock sync.Mutex
+	waiters     map[watchKey]*nats.Subscription
+}
+
+// Shortcut reports:
+// When the UI connects a WS to the query service, a goroutine periodically
+// published rendered reports to that ws.  This process can be interrupted by
+// "shortcut" reports, causing the query service to push a render report
+// immediately. This whole process is controlled by the aforementioned
+// goroutine registering a channel with the collector.  We store these
+// registered channels in a map keyed by the userid and the channel itself,
+// which in go is hashable.  We then listen on a NATS topic for any shortcut
+// reports coming from the collection service.
+type watchKey struct {
+	userid string
+	c      chan struct{}
 }
 
 // NewDynamoDBCollector the reaper of souls
 // https://github.com/aws/aws-sdk-go/wiki/common-examples
-func NewDynamoDBCollector(dynamoDBConfig, s3Config *aws.Config, userIDer UserIDer, tableName, bucketName string) DynamoDBCollector {
+func NewDynamoDBCollector(
+	userIDer UserIDer,
+	dynamoDBConfig, s3Config *aws.Config,
+	tableName, bucketName, natsHost string,
+) (DynamoDBCollector, error) {
+	var nc *nats.Conn
+	if natsHost != "" {
+		var err error
+		nc, err = nats.Connect(natsHost)
+		if err != nil {
+			return nil, err
+		}
+	}
+
 	return &dynamoDBCollector{
 		db:         dynamodb.New(session.New(dynamoDBConfig)),
 		s3:         s3.New(session.New(s3Config)),
 		userIDer:   userIDer,
 		tableName:  tableName,
 		bucketName: bucketName,
 		merger:     app.NewSmartMerger(),
-		cache:      gcache.New(cacheSize).LRU().Expiration(cacheExpiration).Build(),
-	}
+		cache:      gcache.New(reportCacheSize).LRU().Expiration(reportCacheExpiration).Build(),
+		nats:       nc,
+		waiters:    map[watchKey]*nats.Subscription{},
+	}, nil
 }
 
 // CreateDynamoDBTables creates the required tables in dynamodb
@@ -163,21 +206,6 @@ func (c *dynamoDBCollector) CreateTables() error {
 	return err
 }
 
-func errorCode(err error) string {
-	if err == nil {
-		return "200"
-	}
-	return "500"
-}
-
-func timeRequest(method string, metric *prometheus.SummaryVec, f func() error) error {
-	startTime := time.Now()
-	err := f()
-	duration := time.Now().Sub(startTime)
-	metric.WithLabelValues(method, errorCode(err)).Observe(float64(duration.Nanoseconds()))
-	return err
-}
-
 // getReportKeys gets the s3 keys for reports in this range
 func (c *dynamoDBCollector) getReportKeys(rowKey string, start, end time.Time) ([]string, error) {
 	var resp *dynamodb.QueryOutput
@@ -419,9 +447,82 @@ func (c *dynamoDBCollector) Add(ctx context.Context, rep report.Report) error {
 		dynamoConsumedCapacity.WithLabelValues("PutItem").
 			Add(float64(*resp.ConsumedCapacity.CapacityUnits))
 	}
-	return err
+	if err != nil {
+		return err
+	}
+
+	if rep.Shortcut && c.nats != nil {
+		err := c.nats.Publish(userid, []byte(s3Key))
+		natsRequests.WithLabelValues("Publish", errorCode(err)).Add(1)
+		if err != nil {
+			log.Errorf("Error sending shortcut report: %v", err)
+		}
+	}
+
+	return nil
+}
+
+func (c *dynamoDBCollector) WaitOn(ctx context.Context, waiter chan struct{}) {
+	userid, err := c.userIDer(ctx)
+	if err != nil {
+		log.Errorf("Error getting user id in WaitOn: %v", err)
+		return
+	}
+
+	if c.nats == nil {
+		return
+	}
+
+	sub, err := c.nats.SubscribeSync(userid)
+	natsRequests.WithLabelValues("SubscribeSync", errorCode(err)).Add(1)
+	if err != nil {
+		log.Errorf("Error subscribing for shortcuts: %v", err)
+		return
+	}
+
+	c.waitersLock.Lock()
+	c.waiters[watchKey{userid, waiter}] = sub
+	c.waitersLock.Unlock()
+
+	go func() {
+		for {
+			_, err := sub.NextMsg(natsTimeout)
+			if err == nats.ErrTimeout {
+				continue
+			}
+			natsRequests.WithLabelValues("NextMsg", errorCode(err)).Add(1)
+			if err != nil {
+				log.Debugf("NextMsg error: %v", err)
+				return
+			}
+			select {
+			case waiter <- struct{}{}:
+			default:
+			}
+		}
+	}()
 }
 
-func (c *dynamoDBCollector) WaitOn(context.Context, chan struct{}) {}
+func (c *dynamoDBCollector) UnWait(ctx context.Context, waiter chan struct{}) {
+	userid, err := c.userIDer(ctx)
+	if err != nil {
+		log.Errorf("Error getting user id in WaitOn: %v", err)
+		return
+	}
+
+	if c.nats == nil {
+		return
+	}
+
+	c.waitersLock.Lock()
+	key := watchKey{userid, waiter}
+	sub := c.waiters[key]
+	delete(c.waiters, key)
+	c.waitersLock.Unlock()
 
-func (c *dynamoDBCollector) UnWait(context.Context, chan struct{}) {}
+	err = sub.Unsubscribe()
+	natsRequests.WithLabelValues("Unsubscribe", errorCode(err)).Add(1)
+	if err != nil {
+		log.Errorf("Error on unsubscribe: %v", err)
+	}
+}

prog/app.go

@@ -76,7 +76,7 @@ func awsConfigFromURL(url *url.URL) (*aws.Config, error) {
 	return config, nil
 }
 
-func collectorFactory(userIDer multitenant.UserIDer, collectorURL, s3URL string, window time.Duration, createTables bool) (app.Collector, error) {
+func collectorFactory(userIDer multitenant.UserIDer, collectorURL, s3URL, natsHostname string, window time.Duration, createTables bool) (app.Collector, error) {
 	if collectorURL == "local" {
 		return app.NewCollector(window), nil
 	}
@@ -101,8 +101,11 @@ func collectorFactory(userIDer multitenant.UserIDer, collectorURL, s3URL string,
 		}
 		tableName := strings.TrimPrefix(parsed.Path, "/")
 		bucketName := strings.TrimPrefix(s3.Path, "/")
-		dynamoCollector := multitenant.NewDynamoDBCollector(
-			dynamoDBConfig, s3Config, userIDer, tableName, bucketName)
+		dynamoCollector, err := multitenant.NewDynamoDBCollector(
+			userIDer, dynamoDBConfig, s3Config, tableName, bucketName, natsHostname)
+		if err != nil {
+			return nil, err
+		}
 		if createTables {
 			if err := dynamoCollector.CreateTables(); err != nil {
 				return nil, err
@@ -167,12 +170,20 @@ func appMain(flags appFlags) {
 	setLogLevel(flags.logLevel)
 	setLogFormatter(flags.logPrefix)
 
+	defer log.Info("app exiting")
+	rand.Seed(time.Now().UnixNano())
+	app.UniqueID = strconv.FormatInt(rand.Int63(), 16)
+	app.Version = version
+	log.Infof("app starting, version %s, ID %s", app.Version, app.UniqueID)
+	log.Infof("command line: %v", os.Args)
+
 	userIDer := multitenant.NoopUserIDer
 	if flags.userIDHeader != "" {
 		userIDer = multitenant.UserIDHeader(flags.userIDHeader)
 	}
 
-	collector, err := collectorFactory(userIDer, flags.collectorURL, flags.s3URL, flags.window, flags.awsCreateTables)
+	collector, err := collectorFactory(
+		userIDer, flags.collectorURL, flags.s3URL, flags.natsHostname, flags.window, flags.awsCreateTables)
 	if err != nil {
 		log.Fatalf("Error creating collector: %v", err)
 		return
@@ -190,13 +201,6 @@ func appMain(flags appFlags) {
 		return
 	}
 
-	defer log.Info("app exiting")
-	rand.Seed(time.Now().UnixNano())
-	app.UniqueID = strconv.FormatInt(rand.Int63(), 16)
-	app.Version = version
-	log.Infof("app starting, version %s, ID %s", app.Version, app.UniqueID)
-	log.Infof("command line: %v", os.Args)
-
 	// Start background version checking
 	checkpoint.CheckInterval(&checkpoint.CheckParams{
 		Product: "scope-app",

prog/main.go

@@ -100,6 +100,7 @@ type appFlags struct {
 	s3URL            string
 	controlRouterURL string
 	pipeRouterURL    string
+	natsHostname     string
 	userIDHeader     string
 
 	awsCreateTables bool
@@ -176,6 +177,7 @@ func main() {
 	flag.StringVar(&flags.app.s3URL, "app.collector.s3", "local", "S3 URL to use (when collector is dynamodb)")
 	flag.StringVar(&flags.app.controlRouterURL, "app.control.router", "local", "Control router to use (local or sqs)")
 	flag.StringVar(&flags.app.pipeRouterURL, "app.pipe.router", "local", "Pipe router to use (local)")
+	flag.StringVar(&flags.app.natsHostname, "app.nats", "", "Hostname for NATS service to use for shortcut reports.  If empty, shortcut reporting will be disabled.")
 	flag.StringVar(&flags.app.userIDHeader, "app.userid.header", "", "HTTP header to use as userid")
 
 	flag.BoolVar(&flags.app.awsCreateTables, "app.aws.create.tables", false, "Create the tables in DynamoDB")