Address pull request comments

* Use RWMutex instead of Mutex
* Use time.Duration instead of int for unlock time
* Use time.After with select instead of time.Sleep

2 files changed, 19 insertions, 17 deletions

diff --git a/accounts/account_manager.go b/accounts/account_manager.go
index 4d63bd0f2..38dd6f736 100644
--- a/accounts/account_manager.go
+++ b/accounts/account_manager.go
@@ -48,19 +48,19 @@ type Account struct {
 }
 
 type AccountManager struct {
-	keyStore             crypto.KeyStore2
-	unlockedKeys         map[string]crypto.Key
-	unlockedMilliSeconds int
-	mutex                sync.Mutex
+	keyStore           crypto.KeyStore2
+	unlockedKeys       map[string]crypto.Key
+	unlockMilliseconds time.Duration
+	mutex              sync.RWMutex
 }
 
-func NewAccountManager(keyStore crypto.KeyStore2, unlockMilliSeconds int) AccountManager {
+func NewAccountManager(keyStore crypto.KeyStore2, unlockMilliseconds time.Duration) AccountManager {
 	keysMap := make(map[string]crypto.Key)
 	am := &AccountManager{
-		keyStore:             keyStore,
-		unlockedKeys:         keysMap,
-		unlockedMilliSeconds: unlockMilliSeconds,
-		mutex:                sync.Mutex{}, // for accessing unlockedKeys map
+		keyStore:           keyStore,
+		unlockedKeys:       keysMap,
+		unlockMilliseconds: unlockMilliseconds,
+		mutex:              sync.RWMutex{}, // for accessing unlockedKeys map
 	}
 	return *am
 }
@@ -70,9 +70,9 @@ func (am AccountManager) DeleteAccount(address []byte, auth string) error {
 }
 
 func (am *AccountManager) Sign(fromAccount *Account, toSign []byte) (signature []byte, err error) {
-	am.mutex.Lock()
+	am.mutex.RLock()
 	unlockedKey := am.unlockedKeys[string(fromAccount.Address)]
-	am.mutex.Unlock()
+	am.mutex.RUnlock()
 	if unlockedKey.Address == nil {
 		return nil, ErrLocked
 	}
@@ -85,9 +85,9 @@ func (am *AccountManager) SignLocked(fromAccount *Account, keyAuth string, toSig
 	if err != nil {
 		return nil, err
 	}
-	am.mutex.Lock()
+	am.mutex.RLock()
 	am.unlockedKeys[string(fromAccount.Address)] = *key
-	am.mutex.Unlock()
+	am.mutex.RUnlock()
 	go unlockLater(am, fromAccount.Address)
 	signature, err = crypto.Sign(toSign, key.PrivateKey)
 	return signature, err
@@ -121,9 +121,11 @@ func (am *AccountManager) Accounts() ([]Account, error) {
 }
 
 func unlockLater(am *AccountManager, addr []byte) {
-	time.Sleep(time.Millisecond * time.Duration(am.unlockedMilliSeconds))
-	am.mutex.Lock()
+	select {
+	case <-time.After(time.Millisecond * am.unlockMilliseconds):
+	}
+	am.mutex.RLock()
 	// TODO: how do we know the key is actually gone from memory?
 	delete(am.unlockedKeys, string(addr))
-	am.mutex.Unlock()
+	am.mutex.RUnlock()
 }
diff --git a/accounts/accounts_test.go b/accounts/accounts_test.go
index 8f036fd1f..44d1d72f1 100644
--- a/accounts/accounts_test.go
+++ b/accounts/accounts_test.go
@@ -21,7 +21,7 @@ func TestAccountManager(t *testing.T) {
 	}
 
 	// Cleanup
-	time.Sleep(time.Millisecond * time.Duration(150)) // wait for locking
+	time.Sleep(time.Millisecond * 150) // wait for locking
 
 	accounts, err := am.Accounts()
 	if err != nil {