services/horizon/ledger: Fix memory leak in HistoryDBSource

services/horizon/CHANGELOG.md

@@ -7,6 +7,7 @@ file. This project adheres to [Semantic Versioning](http://semver.org/).
 
 * Add `last_modified_time` to account responses. `last_modified_time` is the
  closing time of the most recent ledger in which the account was modified.
+* Fix a memory leak in the code responsible for streaming [#2548](https://github.com/stellar/go/pull/2548).
 
 ## v1.2.1
 

services/horizon/internal/ledger/ledger_source.go

@@ -6,10 +6,12 @@ import (
 )
 
 // Source exposes two helpers methods to help you find out the current
-// ledger and yield every time there is a new ledger.
+// ledger and yield every time there is a new ledger. Call `Close` when
+// source is no longer used.
 type Source interface {
 	CurrentLedger() uint32
 	NextLedger(currentSequence uint32) chan uint32
+	Close()
 }
 
 type currentStateFunc func() State
@@ -19,23 +21,24 @@ type currentStateFunc func() State
 type HistoryDBSource struct {
 	updateFrequency time.Duration
 	currentState    currentStateFunc
+	closed          bool
 }
 
 // NewHistoryDBSource constructs a new instance of HistoryDBSource
-func NewHistoryDBSource(updateFrequency time.Duration) HistoryDBSource {
-	return HistoryDBSource{
+func NewHistoryDBSource(updateFrequency time.Duration) *HistoryDBSource {
+	return &HistoryDBSource{
 		updateFrequency: updateFrequency,
 		currentState:    CurrentState,
 	}
 }
 
 // CurrentLedger returns the current ledger.
-func (source HistoryDBSource) CurrentLedger() uint32 {
+func (source *HistoryDBSource) CurrentLedger() uint32 {
 	return source.currentState().ExpHistoryLatest
 }
 
 // NextLedger returns a channel which yields every time there is a new ledger with a sequence number larger than currentSequence.
-func (source HistoryDBSource) NextLedger(currentSequence uint32) chan uint32 {
+func (source *HistoryDBSource) NextLedger(currentSequence uint32) chan uint32 {
 	// Make sure this is buffered channel of size 1. Otherwise, the go routine below
 	// will never return if `newLedgers` channel is not read. From Effective Go:
 	// > If the channel is unbuffered, the sender blocks until the receiver has received the value.
@@ -46,6 +49,10 @@ func (source HistoryDBSource) NextLedger(currentSequence uint32) chan uint32 {
 				time.Sleep(source.updateFrequency)
 			}
 
+			if source.closed {
+				return
+			}
+
 			currentLedgerState := source.currentState()
 			if currentLedgerState.ExpHistoryLatest > currentSequence {
 				newLedgers <- currentLedgerState.ExpHistoryLatest
@@ -57,6 +64,11 @@ func (source HistoryDBSource) NextLedger(currentSequence uint32) chan uint32 {
 	return newLedgers
 }
 
+// Close closes the internal go routines.
+func (source *HistoryDBSource) Close() {
+	source.closed = true
+}
+
 // TestingSource is helper struct which implements the LedgerSource
 // interface.
 type TestingSource struct {
@@ -106,3 +118,5 @@ func (source *TestingSource) NextLedger(currentSequence uint32) chan uint32 {
 
 	return response
 }
+
+func (source *TestingSource) Close() {}

services/horizon/internal/render/sse/stream_handler.go

@@ -8,10 +8,14 @@ import (
 	"github.com/stellar/throttled"
 )
 
+type LedgerSourceFactory interface {
+	Get() ledger.Source
+}
+
 // StreamHandler represents a stream handling action
 type StreamHandler struct {
-	RateLimiter  *throttled.HTTPRateLimiter
-	LedgerSource ledger.Source
+	RateLimiter         *throttled.HTTPRateLimiter
+	LedgerSourceFactory LedgerSourceFactory
 }
 
 // GenerateEventsFunc generates a slice of sse.Event which are sent via
@@ -30,7 +34,10 @@ func (handler StreamHandler) ServeStream(
 	stream := NewStream(ctx, w)
 	stream.SetLimit(limit)
 
-	currentLedgerSequence := handler.LedgerSource.CurrentLedger()
+	ledgerSource := handler.LedgerSourceFactory.Get()
+	defer ledgerSource.Close()
+
+	currentLedgerSequence := ledgerSource.CurrentLedger()
 	for {
 		// Rate limit the request if it's a call to stream since it queries the DB every second. See
 		// https://github.com/stellar/go/issues/715 for more details.
@@ -71,7 +78,7 @@ func (handler StreamHandler) ServeStream(
 		stream.Init()
 
 		select {
-		case currentLedgerSequence = <-handler.LedgerSource.NextLedger(currentLedgerSequence):
+		case currentLedgerSequence = <-ledgerSource.NextLedger(currentLedgerSequence):
 			continue
 		case <-ctx.Done():
 			stream.Done()

services/horizon/internal/render/sse/stream_handler_test.go

@@ -9,11 +9,17 @@ import (
 	"github.com/stellar/go/services/horizon/internal/ledger"
 )
 
+type testingFactory struct {
+	ledgerSource ledger.Source
+}
+
+func (f *testingFactory) Get() ledger.Source {
+	return f.ledgerSource
+}
+
 func TestSendByeByeOnContextDone(t *testing.T) {
 	ledgerSource := ledger.NewTestingSource(1)
-	handler := StreamHandler{
-		LedgerSource: ledgerSource,
-	}
+	handler := StreamHandler{LedgerSourceFactory: &testingFactory{ledgerSource}}
 
 	r, err := http.NewRequest("GET", "http://localhost", nil)
 	if err != nil {

services/horizon/internal/stream_handler_test.go

@@ -21,6 +21,14 @@ import (
 	"github.com/stellar/go/support/render/hal"
 )
 
+type testingFactory struct {
+	ledgerSource ledger.Source
+}
+
+func (f *testingFactory) Get() ledger.Source {
+	return f.ledgerSource
+}
+
 // StreamTest utility struct to wrap SSE related tests.
 type StreamTest struct {
 	ledgerSource  *ledger.TestingSource
@@ -64,7 +72,7 @@ func NewStreamablePageTest(
 ) *StreamTest {
 	ledgerSource := ledger.NewTestingSource(currentLedger)
 	action.ledgerSource = ledgerSource
-	streamHandler := sse.StreamHandler{LedgerSource: ledgerSource}
+	streamHandler := sse.StreamHandler{LedgerSourceFactory: &testingFactory{ledgerSource}}
 	handler := streamablePageHandler(action, streamHandler)
 
 	return newStreamTest(
@@ -85,7 +93,7 @@ func NewStreamableObjectTest(
 ) *StreamTest {
 	ledgerSource := ledger.NewTestingSource(currentLedger)
 	action.ledgerSource = ledgerSource
-	streamHandler := sse.StreamHandler{LedgerSource: ledgerSource}
+	streamHandler := sse.StreamHandler{LedgerSourceFactory: &testingFactory{ledgerSource}}
 	handler := streamableObjectActionHandler{action: action, limit: limit, streamHandler: streamHandler}
 
 	return newStreamTest(

services/horizon/internal/web.go

@@ -131,6 +131,14 @@ func (w *web) mustInstallMiddlewares(app *App, connTimeout time.Duration) {
 	w.internalRouter.Use(loggerMiddleware)
 }
 
+type historyLedgerSourceFactory struct {
+	updateFrequency time.Duration
+}
+
+func (f historyLedgerSourceFactory) Get() ledger.Source {
+	return ledger.NewHistoryDBSource(f.updateFrequency)
+}
+
 // mustInstallActions installs the routing configuration of horizon onto the
 // provided app.  All route registration should be implemented here.
 func (w *web) mustInstallActions(
@@ -151,8 +159,8 @@ func (w *web) mustInstallActions(
 	r.Get("/", RootAction{}.Handle)
 
 	streamHandler := sse.StreamHandler{
-		RateLimiter:  w.rateLimiter,
-		LedgerSource: ledger.NewHistoryDBSource(w.sseUpdateFrequency),
+		RateLimiter:         w.rateLimiter,
+		LedgerSourceFactory: historyLedgerSourceFactory{updateFrequency: w.sseUpdateFrequency},
 	}
 
 	// State endpoints behind stateMiddleware