Use rate/limit, make Writer a method

coadler · coadler · commit b2ecc19f33e7 · 2020-06-15T16:40:44.000-05:00
diff --git a/cmd/coder/shell.go b/cmd/coder/shell.go
@@ -147,10 +147,8 @@ func runCommand(ctx context.Context, envName string, command string, args []stri
 		stdin := process.Stdin()
 		defer stdin.Close()
 
-		ap := &activity.Pusher{Source: sshActivityName, EnvID: env.ID, Client: entClient}
-		defer ap.Start()()
-
-		wr := activity.Writer(ap, stdin)
+		ap := activity.NewPusher(entClient, env.ID, sshActivityName)
+		wr := ap.Writer(stdin)
 		_, err := io.Copy(wr, os.Stdin)
 		if err != nil {
 			cancel()
diff --git a/internal/activity/pusher.go b/internal/activity/pusher.go
@@ -1,55 +1,41 @@
 package activity
 
 import (
-	"context"
-	"sync/atomic"
 	"time"
 
 	"cdr.dev/coder-cli/internal/entclient"
 	"go.coder.com/flog"
+	"golang.org/x/time/rate"
 )
 
 const pushInterval = time.Minute
 
 // Pusher pushes activity metrics no more than once per pushInterval. Pushes
 // within the same interval are a no-op.
 type Pusher struct {
-	Source string
-	EnvID  string
-	Client *entclient.Client
+	envID  string
+	source string
 
-	state int64
+	client *entclient.Client
+	rate   *rate.Limiter
 }
 
-func (p *Pusher) Push() {
-	if atomic.CompareAndSwapInt64(&p.state, 0, 1) {
-		err := p.Client.PushActivity(p.Source, p.EnvID)
-		if err != nil {
-			flog.Error("failed to push activity: %s", err.Error())
-		}
+func NewPusher(c *entclient.Client, envID, source string) *Pusher {
+	return &Pusher{
+		envID:  envID,
+		source: source,
+		client: c,
+		rate:   rate.NewLimiter(rate.Every(pushInterval), 1),
 	}
 }
 
-// Start starts the reset routine. It resets the state every pushInterval,
-// allowing Push to push a new activity.
-func (p *Pusher) Start() (end func()) {
-	ctx, cancel := context.WithCancel(context.Background())
-	go func() {
-		tick := time.NewTicker(pushInterval)
-		defer tick.Stop()
-
-		for {
-			select {
-			case <-ctx.Done():
-				return
-			case <-tick.C:
-				atomic.StoreInt64(&p.state, 0)
-			}
-
-		}
-	}()
-
-	return func() {
-		cancel()
+func (p *Pusher) Push() {
+	if !p.rate.Allow() {
+		return
+	}
+
+	err := p.client.PushActivity(p.source, p.envID)
+	if err != nil {
+		flog.Error("push activity: %s", err.Error())
 	}
 }
diff --git a/internal/activity/writer.go b/internal/activity/writer.go
@@ -12,6 +12,6 @@ func (w *activityWriter) Write(p []byte) (n int, err error) {
 	return w.wr.Write(p)
 }
 
-func Writer(p *Pusher, wr io.Writer) io.Writer {
+func (p *Pusher) Writer(wr io.Writer) io.Writer {
 	return &activityWriter{p: p, wr: wr}
 }
diff --git a/internal/sync/sync.go b/internal/sync/sync.go
@@ -260,7 +260,7 @@ func (s Sync) Run() error {
 	}
 	defer notify.Stop(events)
 
-	ap := activity.Pusher{Source: activityName, EnvID: s.Env.ID, Client: s.Client}
+	rp := activity.Pusher{Source: activityName, EnvID: s.Env.ID, Client: s.Client}
 	defer ap.Start()()
 	ap.Push()
 

Original file line number	Diff line number	Diff line change
`@@ -12,6 +12,6 @@ func (w *activityWriter) Write(p []byte) (n int, err error) {`
`12`	`12`	`return w.wr.Write(p)`
`13`	`13`	`}`
`14`	`14`
`15`		`-func Writer(p *Pusher, wr io.Writer) io.Writer {`
	`15`	`+func (p *Pusher) Writer(wr io.Writer) io.Writer {`
`16`	`16`	`return &activityWriter{p: p, wr: wr}`
`17`	`17`	`}`
Original file line number	Diff line number	Diff line change
`@@ -260,7 +260,7 @@ func (s Sync) Run() error {`
`260`	`260`	`}`
`261`	`261`	`defer notify.Stop(events)`
`262`	`262`
`263`		`- ap := activity.Pusher{Source: activityName, EnvID: s.Env.ID, Client: s.Client}`
	`263`	`+ rp := activity.Pusher{Source: activityName, EnvID: s.Env.ID, Client: s.Client}`
`264`	`264`	`defer ap.Start()()`
`265`	`265`	`ap.Push()`
`266`	`266`