From 36093f60baed21c8299fbb31f26c6e9123725fda Mon Sep 17 00:00:00 2001
From: Zeal Jagannatha <zealjagannatha@gmail.com>
Date: Tue, 24 May 2016 12:11:35 -0700
Subject: [PATCH] Fixed #270

---
 src/daemon/atcd.go      | 9 +++++++--
 src/daemon/atcd_test.go | 6 ++----
 2 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/src/daemon/atcd.go b/src/daemon/atcd.go
index e2ddd6a..64723ba 100644
--- a/src/daemon/atcd.go
+++ b/src/daemon/atcd.go
@@ -136,6 +136,7 @@ func (atcd *Atcd) CreateGroup(member string) (*atc_thrift.ShapingGroup, error) {
 		return nil, err
 	}
 	grp.ID = dbgrp.id
+	defer atcd.Cleanup()
 	return grp, nil
 }
 
@@ -222,6 +223,7 @@ func (atcd *Atcd) JoinGroup(id int64, to_add, token string) error {
 	if err := atcd.shaper.JoinGroup(id, tgt); err != nil {
 		return err
 	}
+	defer atcd.Cleanup()
 	return err
 }
 
@@ -250,8 +252,7 @@ func (atcd *Atcd) LeaveGroup(id int64, to_remove, token string) error {
 	if err := atcd.shaper.LeaveGroup(id, tgt); err != nil {
 		return err
 	}
-	// FIXME: clean shaper's group too!
-	defer atcd.db.Cleanup()
+	defer atcd.Cleanup()
 	return atcd.db.DeleteMember(tgt)
 }
 
@@ -318,6 +319,10 @@ func (atcd *Atcd) otp(group *DbGroup) *otp.TOTP {
 	}
 }
 
+func (atcd *Atcd) Cleanup() {
+	atcd.db.Cleanup()
+}
+
 func makeSecret() string {
 	// Can probably find a better source of random secrets than this...
 	return uuid.New()
diff --git a/src/daemon/atcd_test.go b/src/daemon/atcd_test.go
index bf5d56f..1ae3f6d 100644
--- a/src/daemon/atcd_test.go
+++ b/src/daemon/atcd_test.go
@@ -1,8 +1,8 @@
 package daemon
 
 import (
-	"runtime"
 	"testing"
+	"time"
 
 	"github.com/facebook/augmented-traffic-control/src/atc_thrift"
 	"github.com/facebook/augmented-traffic-control/src/shaping"
@@ -46,15 +46,13 @@ func TestAtcdLeaveGroup(_t *testing.T) {
 }
 
 func TestAtcdCleansEmptyGroups(_t *testing.T) {
-	_t.Skip("something changed with how goroutines were run and this is broken in 1.6. Deferring fixing until we refactor the cleanup logic")
 	t := Setup(_t, true)
 	defer t.Cleanup()
 	grp := t.CreateGroup("1.2.3.4")
 	token := t.token(grp)
 	t.LeaveGroup(grp.ID, "1.2.3.4", token)
 
-	// Allow other goroutines to run
-	runtime.Gosched()
+	time.Sleep(1 * time.Millisecond)
 
 	grp, err := t.atcd.GetGroup(grp.ID)
 	if err != NoSuchItem {