persistence: expose close methods to cleanup gocql session conns (#240)

Author: venkat1109

common/mocks/ExecutionManager.go

@@ -28,6 +28,11 @@ type ExecutionManager struct {
 	mock.Mock
 }
 
+// Close provides a mock function with given fields:
+func (_m *ExecutionManager) Close() {
+	_m.Called()
+}
+
 // CompleteTimerTask provides a mock function with given fields: request
 func (_m *ExecutionManager) CompleteTimerTask(request *persistence.CompleteTimerTaskRequest) error {
 	ret := _m.Called(request)

common/mocks/HistoryManager.go

@@ -28,6 +28,11 @@ type HistoryManager struct {
 	mock.Mock
 }
 
+// Close provides a mock function with given fields:
+func (_m *HistoryManager) Close() {
+	_m.Called()
+}
+
 // AppendHistoryEvents provides a mock function with given fields: request
 func (_m *HistoryManager) AppendHistoryEvents(request *persistence.AppendHistoryEventsRequest) error {
 	ret := _m.Called(request)

common/mocks/MetadataManager.go

@@ -28,6 +28,11 @@ type MetadataManager struct {
 	mock.Mock
 }
 
+// Close provides a mock function with given fields:
+func (_m *MetadataManager) Close() {
+	_m.Called()
+}
+
 // CreateDomain provides a mock function with given fields: request
 func (_m *MetadataManager) CreateDomain(request *persistence.CreateDomainRequest) (*persistence.CreateDomainResponse, error) {
 	ret := _m.Called(request)

common/mocks/ShardManager.go

@@ -28,6 +28,11 @@ type ShardManager struct {
 	mock.Mock
 }
 
+// Close provides a mock function with given fields:
+func (_m *ShardManager) Close() {
+	_m.Called()
+}
+
 // CreateShard provides a mock function with given fields: request
 func (_m *ShardManager) CreateShard(request *persistence.CreateShardRequest) error {
 	ret := _m.Called(request)

common/mocks/TaskManager.go

@@ -28,6 +28,11 @@ type TaskManager struct {
 	mock.Mock
 }
 
+// Close provides a mock function with given fields:
+func (_m *TaskManager) Close() {
+	_m.Called()
+}
+
 // LeaseTaskList provides a mock function with given fields: request
 func (_m *TaskManager) LeaseTaskList(request *persistence.LeaseTaskListRequest) (*persistence.LeaseTaskListResponse, error) {
 	ret := _m.Called(request)

common/mocks/VisibilityManager.go

@@ -28,6 +28,11 @@ type VisibilityManager struct {
 	mock.Mock
 }
 
+// Close provides a mock function with given fields:
+func (_m *VisibilityManager) Close() {
+	_m.Called()
+}
+
 // GetClosedWorkflowExecution provides a mock function with given fields: request
 func (_m *VisibilityManager) GetClosedWorkflowExecution(request *persistence.GetClosedWorkflowExecutionRequest) (*persistence.GetClosedWorkflowExecutionResponse, error) {
 	ret := _m.Called(request)

common/persistence/cassandraHistoryPersistence.go

@@ -77,6 +77,13 @@ func NewCassandraHistoryPersistence(hosts string, dc string, keyspace string, lo
 	return &cassandraHistoryPersistence{session: session, logger: logger}, nil
 }
 
+// Close gracefully releases the resources held by this object
+func (h *cassandraHistoryPersistence) Close() {
+	if h.session != nil {
+		h.session.Close()
+	}
+}
+
 func (h *cassandraHistoryPersistence) AppendHistoryEvents(request *AppendHistoryEventsRequest) error {
 	var query *gocql.Query
 	if request.Overwrite {

common/persistence/cassandraMetadataPersistence.go

@@ -105,6 +105,13 @@ func NewCassandraMetadataPersistence(hosts string, dc string, keyspace string, l
 	return &cassandraMetadataPersistence{session: session, logger: logger}, nil
 }
 
+// Close releases the resources held by this object
+func (m *cassandraMetadataPersistence) Close() {
+	if m.session != nil {
+		m.session.Close()
+	}
+}
+
 // Cassandra does not support conditional updates across multiple tables.  For this reason we have to first insert into
 // 'Domains' table and then do a conditional insert into domains_by_name table.  If the conditional write fails we
 // delete the orphaned entry from domains table.  There is a chance delete entry could fail and we never delete the

common/persistence/cassandraPersistence.go

@@ -517,6 +517,13 @@ func NewCassandraTaskPersistence(hosts string, dc string, keyspace string, logge
 	return &cassandraPersistence{shardID: -1, session: session, lowConslevel: gocql.One, logger: logger}, nil
 }
 
+// Close releases the underlying resources held by this object
+func (d *cassandraPersistence) Close() {
+	if d.session != nil {
+		d.session.Close()
+	}
+}
+
 func (d *cassandraPersistence) CreateShard(request *CreateShardRequest) error {
 	cqlNowTimestamp := common.UnixNanoToCQLTimestamp(time.Now().UnixNano())
 	shardInfo := request.ShardInfo

common/persistence/cassandraVisibilityPersistence.go

@@ -140,6 +140,13 @@ func NewCassandraVisibilityPersistence(
 	return &cassandraVisibilityPersistence{session: session, lowConslevel: gocql.One, logger: logger}, nil
 }
 
+// Close releases the resources held by this object
+func (v *cassandraVisibilityPersistence) Close() {
+	if v.session != nil {
+		v.session.Close()
+	}
+}
+
 func (v *cassandraVisibilityPersistence) RecordWorkflowExecutionStarted(
 	request *RecordWorkflowExecutionStartedRequest) error {
 	query := v.session.Query(templateCreateWorkflowExecutionStarted,

common/persistence/dataInterfaces.go

@@ -594,15 +594,22 @@ type (
 		Name string
 	}
 
+	// Closeable is an interface for any entity that supports a close operation to release resources
+	Closeable interface {
+		Close()
+	}
+
 	// ShardManager is used to manage all shards
 	ShardManager interface {
+		Closeable
 		CreateShard(request *CreateShardRequest) error
 		GetShard(request *GetShardRequest) (*GetShardResponse, error)
 		UpdateShard(request *UpdateShardRequest) error
 	}
 
 	// ExecutionManager is used to manage workflow executions
 	ExecutionManager interface {
+		Closeable
 		CreateWorkflowExecution(request *CreateWorkflowExecutionRequest) (*CreateWorkflowExecutionResponse, error)
 		GetWorkflowExecution(request *GetWorkflowExecutionRequest) (*GetWorkflowExecutionResponse, error)
 		UpdateWorkflowExecution(request *UpdateWorkflowExecutionRequest) error
@@ -623,6 +630,7 @@ type (
 
 	// TaskManager is used to manage tasks
 	TaskManager interface {
+		Closeable
 		LeaseTaskList(request *LeaseTaskListRequest) (*LeaseTaskListResponse, error)
 		UpdateTaskList(request *UpdateTaskListRequest) (*UpdateTaskListResponse, error)
 		CreateTasks(request *CreateTasksRequest) (*CreateTasksResponse, error)
@@ -632,6 +640,7 @@ type (
 
 	// HistoryManager is used to manage Workflow Execution HistoryEventBatch
 	HistoryManager interface {
+		Closeable
 		AppendHistoryEvents(request *AppendHistoryEventsRequest) error
 		// GetWorkflowExecutionHistory retrieves the paginated list of history events for given execution
 		GetWorkflowExecutionHistory(request *GetWorkflowExecutionHistoryRequest) (*GetWorkflowExecutionHistoryResponse,
@@ -641,6 +650,7 @@ type (
 
 	// MetadataManager is used to manage metadata CRUD for various entities
 	MetadataManager interface {
+		Closeable
 		CreateDomain(request *CreateDomainRequest) (*CreateDomainResponse, error)
 		GetDomain(request *GetDomainRequest) (*GetDomainResponse, error)
 		UpdateDomain(request *UpdateDomainRequest) error

common/persistence/persistenceMetricClients.go

@@ -154,6 +154,10 @@ func (p *shardPersistenceClient) UpdateShard(request *UpdateShardRequest) error
 	return err
 }
 
+func (p *shardPersistenceClient) Close() {
+	p.persistence.Close()
+}
+
 func (p *workflowExecutionPersistenceClient) CreateWorkflowExecution(request *CreateWorkflowExecutionRequest) (*CreateWorkflowExecutionResponse, error) {
 	p.metricClient.IncCounter(metrics.PersistenceCreateWorkflowExecutionScope, metrics.PersistenceRequests)
 
@@ -296,6 +300,10 @@ func (p *workflowExecutionPersistenceClient) updateErrorMetric(scope int, err er
 	}
 }
 
+func (p *workflowExecutionPersistenceClient) Close() {
+	p.persistence.Close()
+}
+
 func (p *taskPersistenceClient) CreateTasks(request *CreateTasksRequest) (*CreateTasksResponse, error) {
 	p.metricClient.IncCounter(metrics.PersistenceCreateTaskScope, metrics.PersistenceRequests)
 
@@ -378,6 +386,10 @@ func (p *taskPersistenceClient) updateErrorMetric(scope int, err error) {
 	}
 }
 
+func (p *taskPersistenceClient) Close() {
+	p.persistence.Close()
+}
+
 func (p *historyPersistenceClient) AppendHistoryEvents(request *AppendHistoryEventsRequest) error {
 	p.metricClient.IncCounter(metrics.PersistenceAppendHistoryEventsScope, metrics.PersistenceRequests)
 
@@ -436,6 +448,10 @@ func (p *historyPersistenceClient) updateErrorMetric(scope int, err error) {
 	}
 }
 
+func (p *historyPersistenceClient) Close() {
+	p.persistence.Close()
+}
+
 func (p *metadataPersistenceClient) CreateDomain(request *CreateDomainRequest) (*CreateDomainResponse, error) {
 	p.metricClient.IncCounter(metrics.PersistenceCreateDomainScope, metrics.PersistenceRequests)
 
@@ -506,6 +522,10 @@ func (p *metadataPersistenceClient) DeleteDomainByName(request *DeleteDomainByNa
 	return err
 }
 
+func (p *metadataPersistenceClient) Close() {
+	p.persistence.Close()
+}
+
 func (p *metadataPersistenceClient) updateErrorMetric(scope int, err error) {
 	switch err.(type) {
 	case *workflow.DomainAlreadyExistsError:

common/persistence/visibilityInterfaces.go

@@ -105,6 +105,7 @@ type (
 
 	// VisibilityManager is used to manage the visibility store
 	VisibilityManager interface {
+		Closeable
 		RecordWorkflowExecutionStarted(request *RecordWorkflowExecutionStartedRequest) error
 		RecordWorkflowExecutionClosed(request *RecordWorkflowExecutionClosedRequest) error
 		ListOpenWorkflowExecutions(request *ListWorkflowExecutionsRequest) (*ListWorkflowExecutionsResponse, error)

service/frontend/handler.go

@@ -121,6 +121,9 @@ func (wh *WorkflowHandler) Start(thriftService []thrift.TChanServer) error {
 
 // Stop stops the handler
 func (wh *WorkflowHandler) Stop() {
+	wh.metadataMgr.Close()
+	wh.visibitiltyMgr.Close()
+	wh.historyMgr.Close()
 	wh.Service.Stop()
 }
 

service/history/handler.go

@@ -113,6 +113,10 @@ func (h *Handler) Start(thriftService []thrift.TChanServer) error {
 // Stop stops the handler
 func (h *Handler) Stop() {
 	h.controller.Stop()
+	h.shardManager.Close()
+	h.historyMgr.Close()
+	h.metadataMgr.Close()
+	h.visibilityMgr.Close()
 	h.Service.Stop()
 }
 

service/history/shardController.go

@@ -65,14 +65,14 @@ type (
 	}
 
 	historyShardsItem struct {
-		shardID             int
-		shardMgr            persistence.ShardManager
-		historyMgr          persistence.HistoryManager
-		executionMgrFactory persistence.ExecutionManagerFactory
-		engineFactory       EngineFactory
-		host                *membership.HostInfo
-		logger              bark.Logger
-		metricsClient       metrics.Client
+		shardID       int
+		shardMgr      persistence.ShardManager
+		historyMgr    persistence.HistoryManager
+		executionMgr  persistence.ExecutionManager
+		engineFactory EngineFactory
+		host          *membership.HostInfo
+		logger        bark.Logger
+		metricsClient metrics.Client
 
 		sync.RWMutex
 		engine  Engine
@@ -106,19 +106,25 @@ func newShardController(numberOfShards int, host *membership.HostInfo, resolver
 
 func newHistoryShardsItem(shardID int, shardMgr persistence.ShardManager, historyMgr persistence.HistoryManager,
 	executionMgrFactory persistence.ExecutionManagerFactory, factory EngineFactory, host *membership.HostInfo,
-	logger bark.Logger, reporter metrics.Client) *historyShardsItem {
+	logger bark.Logger, reporter metrics.Client) (*historyShardsItem, error) {
+
+	executionMgr, err := executionMgrFactory.CreateExecutionManager(shardID)
+	if err != nil {
+		return nil, err
+	}
+
 	return &historyShardsItem{
-		shardID:             shardID,
-		shardMgr:            shardMgr,
-		historyMgr:          historyMgr,
-		executionMgrFactory: executionMgrFactory,
-		engineFactory:       factory,
-		host:                host,
+		shardID:       shardID,
+		shardMgr:      shardMgr,
+		historyMgr:    historyMgr,
+		executionMgr:  executionMgr,
+		engineFactory: factory,
+		host:          host,
 		logger: logger.WithFields(bark.Fields{
 			logging.TagHistoryShardID: shardID,
 		}),
 		metricsClient: reporter,
-	}
+	}, nil
 }
 
 func (c *shardController) Start() {
@@ -204,8 +210,11 @@ func (c *shardController) getOrCreateHistoryShardItem(shardID int) (*historyShar
 	}
 
 	if info.Identity() == c.host.Identity() {
-		shardItem := newHistoryShardsItem(shardID, c.shardMgr, c.historyMgr, c.executionMgrFactory, c.engineFactory, c.host,
+		shardItem, err := newHistoryShardsItem(shardID, c.shardMgr, c.historyMgr, c.executionMgrFactory, c.engineFactory, c.host,
 			c.logger, c.metricsClient)
+		if err != nil {
+			return nil, err
+		}
 		c.historyShards[shardID] = shardItem
 		logging.LogShardItemCreatedEvent(shardItem.logger, info.Identity(), shardID)
 		return shardItem, nil
@@ -304,13 +313,8 @@ func (i *historyShardsItem) getOrCreateEngine(shardClosedCh chan<- int) (Engine,
 	}
 
 	logging.LogShardEngineCreatingEvent(i.logger, i.host.Identity(), i.shardID)
-	defer logging.LogShardEngineCreatedEvent(i.logger, i.host.Identity(), i.shardID)
-	executionMgr, err := i.executionMgrFactory.CreateExecutionManager(i.shardID)
-	if err != nil {
-		return nil, err
-	}
 
-	context, err := acquireShard(i.shardID, i.shardMgr, i.historyMgr, executionMgr, i.host.Identity(), shardClosedCh,
+	context, err := acquireShard(i.shardID, i.shardMgr, i.historyMgr, i.executionMgr, i.host.Identity(), shardClosedCh,
 		i.logger, i.metricsClient)
 	if err != nil {
 		return nil, err
@@ -319,18 +323,20 @@ func (i *historyShardsItem) getOrCreateEngine(shardClosedCh chan<- int) (Engine,
 	i.engine = i.engineFactory.CreateEngine(context)
 	i.engine.Start()
 
+	logging.LogShardEngineCreatedEvent(i.logger, i.host.Identity(), i.shardID)
 	return i.engine, nil
 }
 
 func (i *historyShardsItem) stopEngine() {
-	logging.LogShardEngineStoppingEvent(i.logger, i.host.Identity(), i.shardID)
-	defer logging.LogShardEngineStoppedEvent(i.logger, i.host.Identity(), i.shardID)
 	i.Lock()
 	defer i.Unlock()
 
 	if i.engine != nil {
+		logging.LogShardEngineStoppingEvent(i.logger, i.host.Identity(), i.shardID)
 		i.engine.Stop()
 		i.engine = nil
+		i.executionMgr.Close()
+		logging.LogShardEngineStoppedEvent(i.logger, i.host.Identity(), i.shardID)
 	}
 }
 

service/history/shardController_test.go

@@ -440,6 +440,7 @@ func (s *shardControllerSuite) TestShardControllerClosed() {
 func (s *shardControllerSuite) setupMocksForAcquireShard(shardID int, mockEngine *MockHistoryEngine, currentRangeID,
 	newRangeID int64) {
 	mockExecutionMgr := &mmocks.ExecutionManager{}
+	mockExecutionMgr.On("Close").Return()
 	s.mockExecutionMgrFactory.On("CreateExecutionManager", shardID).Return(mockExecutionMgr, nil).Once()
 	mockEngine.On("Start").Return().Once()
 	s.mockServiceResolver.On("Lookup", string(shardID)).Return(s.hostInfo, nil).Twice()

service/matching/handler.go

@@ -69,6 +69,8 @@ func (h *Handler) Start(thriftService []thrift.TChanServer) error {
 
 // Stop stops the handler
 func (h *Handler) Stop() {
+	h.engine.Stop()
+	h.taskPersistence.Close()
 	h.Service.Stop()
 }
 

service/matching/matchingEngineInterfaces.go

@@ -29,6 +29,7 @@ import (
 type (
 	// Engine exposes interfaces for clients to poll for activity and decision tasks.
 	Engine interface {
+		Stop()
 		AddDecisionTask(addRequest *m.AddDecisionTaskRequest) error
 		AddActivityTask(addRequest *m.AddActivityTaskRequest) error
 		PollForDecisionTask(ctx thrift.Context, request *m.PollForDecisionTaskRequest) (*m.PollForDecisionTaskResponse, error)