release checkKey lock correctly (#2703)

Signed-off-by: Tony Lee <tl@hudson-trading.com>

CHANGELOG.md

@@ -7,6 +7,10 @@ Versioning](http://semver.org/spec/v2.0.0.html).
 
 ## Unreleased
 
+### Fixed
+- Fixed a regression in the agent that would not allow proxy checks to be
+run for subsequent executions.
+
 ## [5.2.0] - 2019-02-06
 
 ### Added

agent/check_handler.go

@@ -10,7 +10,7 @@ import (
 	"time"
 
 	"github.com/sensu/sensu-go/agent/transformers"
-	"github.com/sensu/sensu-go/api/core/v2"
+	v2 "github.com/sensu/sensu-go/api/core/v2"
 	"github.com/sensu/sensu-go/asset"
 	"github.com/sensu/sensu-go/command"
 	"github.com/sensu/sensu-go/transport"
@@ -74,15 +74,21 @@ func checkKey(request *v2.CheckRequest) string {
 	return strings.Join(parts, "/")
 }
 
-func (a *Agent) executeCheck(request *v2.CheckRequest, entity *v2.Entity) {
+func (a *Agent) addInProgress(request *v2.CheckRequest) {
 	a.inProgressMu.Lock()
 	a.inProgress[checkKey(request)] = request.Config
 	a.inProgressMu.Unlock()
-	defer func() {
-		a.inProgressMu.Lock()
-		delete(a.inProgress, request.Config.Name)
-		a.inProgressMu.Unlock()
-	}()
+}
+
+func (a *Agent) removeInProgress(request *v2.CheckRequest) {
+	a.inProgressMu.Lock()
+	delete(a.inProgress, checkKey(request))
+	a.inProgressMu.Unlock()
+}
+
+func (a *Agent) executeCheck(request *v2.CheckRequest, entity *v2.Entity) {
+	a.addInProgress(request)
+	defer a.removeInProgress(request)
 
 	checkConfig := request.Config
 	checkAssets := request.Assets

agent/check_handler_internal_test.go

@@ -49,6 +49,65 @@ func TestHandleCheck(t *testing.T) {
 	assert.NoError(agent.handleCheck(payload))
 }
 
+func TestCheckInProgress_GH2704(t *testing.T) {
+	assert := assert.New(t)
+
+	checkConfig := corev2.FixtureCheckConfig("normal-check")
+	request := &corev2.CheckRequest{Config: checkConfig, Issued: time.Now().Unix()}
+	key := checkKey(request)
+	assert.Equal("normal-check", key)
+
+	config, cleanup := FixtureConfig()
+	defer cleanup()
+	agent := NewAgent(config)
+
+	agent.addInProgress(request)
+	agent.inProgressMu.Lock()
+	val, ok := agent.inProgress[key]
+	agent.inProgressMu.Unlock()
+	assert.True(ok)
+	assert.True(agent.checkInProgress(request))
+	assert.Equal(request.Config, val)
+
+	agent.removeInProgress(request)
+	agent.inProgressMu.Lock()
+	val, ok = agent.inProgress[key]
+	agent.inProgressMu.Unlock()
+	assert.False(ok)
+	assert.False(agent.checkInProgress(request))
+	assert.Empty(val)
+}
+
+func TestProxyCheckInProgress_GH2704(t *testing.T) {
+	assert := assert.New(t)
+
+	checkConfig := corev2.FixtureCheckConfig("proxy-check")
+	checkConfig.ProxyEntityName = "proxy-entity"
+	request := &corev2.CheckRequest{Config: checkConfig, Issued: time.Now().Unix()}
+	key := checkKey(request)
+	assert.Equal("proxy-check/proxy-entity", key)
+
+	config, cleanup := FixtureConfig()
+	defer cleanup()
+	agent := NewAgent(config)
+
+	agent.addInProgress(request)
+	agent.inProgressMu.Lock()
+	val, ok := agent.inProgress[key]
+	agent.inProgressMu.Unlock()
+	assert.True(ok)
+	assert.True(agent.checkInProgress(request))
+	assert.Equal(request.Config, val)
+
+	agent.removeInProgress(request)
+	agent.inProgressMu.Lock()
+	val, ok = agent.inProgress[key]
+	agent.inProgressMu.Unlock()
+	assert.False(ok)
+	assert.False(agent.checkInProgress(request))
+	assert.Empty(val)
+}
+
 func TestHandleProxyCheck(t *testing.T) {
 	checkA := corev2.FixtureCheckConfig("check")
 	checkA.ProxyEntityName = "A"