Merge pull request #2103 from anywhy/fix_session_last_price

FIX: [session] fix last price concurrent read and write error

Author: c9s

pkg/bbgo/session.go

@@ -220,6 +220,7 @@ type ExchangeSession struct {
 
 	lastPrices         map[string]fixedpoint.Value
 	lastPriceUpdatedAt time.Time
+	lastPricesMutex    sync.Mutex
 
 	// marketDataStores contains the market data store of each market
 	marketDataStores map[string]*types.MarketDataStore
@@ -571,20 +572,20 @@ func (session *ExchangeSession) Init(ctx context.Context, environ *Environment)
 				session.startPrices[kline.Symbol] = kline.Open
 			}
 
-			session.lastPrices[kline.Symbol] = session.MarketDataStream.(*types.HeikinAshiStream).LastOrigin[kline.Symbol][kline.Interval].Close
+			session.setLastPrice(kline.Symbol, session.MarketDataStream.(*types.HeikinAshiStream).LastOrigin[kline.Symbol][kline.Interval].Close)
 		})
 	} else {
 		session.MarketDataStream.OnKLineClosed(func(kline types.KLine) {
 			if _, ok := session.startPrices[kline.Symbol]; !ok {
-				session.startPrices[kline.Symbol] = kline.Open
+				session.setLastPrice(kline.Symbol, kline.Open)
 			}
 
-			session.lastPrices[kline.Symbol] = kline.Close
+			session.setLastPrice(kline.Symbol, kline.Close)
 		})
 	}
 
 	session.MarketDataStream.OnMarketTrade(func(trade types.Trade) {
-		session.lastPrices[trade.Symbol] = trade.Price
+		session.setLastPrice(trade.Symbol, trade.Price)
 	})
 
 	// session-wide max borrowable updating worker
@@ -738,7 +739,7 @@ func (session *ExchangeSession) initSymbol(ctx context.Context, environ *Environ
 				// update last prices by the given kline
 				lastKLine := kLines[len(kLines)-1]
 				if interval == minInterval {
-					session.lastPrices[symbol] = lastKLine.Close
+					session.setLastPrice(symbol, lastKLine.Close)
 				}
 
 				for _, k := range kLines {
@@ -859,15 +860,21 @@ func (session *ExchangeSession) StartPrice(symbol string) (price fixedpoint.Valu
 }
 
 func (session *ExchangeSession) LastPrice(symbol string) (price fixedpoint.Value, ok bool) {
+	session.lastPricesMutex.Lock()
+	defer session.lastPricesMutex.Unlock()
+
 	price, ok = session.lastPrices[symbol]
 	return price, ok
 }
 
 func (session *ExchangeSession) AllLastPrices() map[string]fixedpoint.Value {
-	return session.lastPrices
+	return session.LastPrices()
 }
 
 func (session *ExchangeSession) LastPrices() map[string]fixedpoint.Value {
+	session.lastPricesMutex.Lock()
+	defer session.lastPricesMutex.Unlock()
+
 	return session.lastPrices
 }
 
@@ -965,12 +972,12 @@ func (session *ExchangeSession) UpdatePrices(ctx context.Context, currencies []s
 		// map things like BTCUSDT = {price}
 		if market, ok := markets[k]; ok {
 			if currency2.IsFiatCurrency(market.BaseCurrency) {
-				session.lastPrices[k] = validPrice.Div(fixedpoint.One)
+				session.setLastPrice(k, validPrice.Div(fixedpoint.One))
 			} else {
-				session.lastPrices[k] = validPrice
+				session.setLastPrice(k, validPrice)
 			}
 		} else {
-			session.lastPrices[k] = v.Last
+			session.setLastPrice(k, v.Last)
 		}
 
 		if v.Time.After(lastTime) {
@@ -1330,3 +1337,10 @@ func (session *ExchangeSession) UpdateMaxBorrowable(ctx context.Context) {
 	}
 	session.marginInfoUpdater.UpdateMaxBorrowable(ctx)
 }
+
+func (session *ExchangeSession) setLastPrice(symbol string, price fixedpoint.Value) {
+	session.lastPricesMutex.Lock()
+	defer session.lastPricesMutex.Unlock()
+
+	session.lastPrices[symbol] = price
+}

pkg/bbgo/session_test.go

@@ -1 +1,48 @@
 package bbgo
+
+import (
+	"sync"
+	"testing"
+
+	"github.com/c9s/bbgo/pkg/fixedpoint"
+)
+
+func TestExchangeSession_LastPricesMutex_ConcurrentAccess(t *testing.T) {
+	session := &ExchangeSession{
+		lastPrices: make(map[string]fixedpoint.Value),
+	}
+
+	var wg sync.WaitGroup
+	symbol := "BTCUSDT"
+	writeCount := 50
+
+	// Concurrently write to lastPrices
+	for i := range writeCount {
+		wg.Add(1)
+		go func(i int) {
+			defer wg.Done()
+			price := fixedpoint.NewFromInt(int64(i))
+			session.setLastPrice(symbol, price)
+		}(i)
+	}
+
+	// Concurrently read from lastPrices
+	for range writeCount {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			_, _ = session.LastPrice(symbol)
+		}()
+	}
+
+	wg.Wait()
+
+	price, ok := session.LastPrice(symbol)
+	if !ok {
+		t.Fatalf("expected price for symbol %s", symbol)
+	}
+
+	if price.Int64() < 0 || price.Int64() >= int64(writeCount) {
+		t.Errorf("unexpected price %d, should be in [0, %d)", price.Int64(), writeCount)
+	}
+}