status-im/status-go

package blockchainstate

import (
    "context"
    "sync"
    "time"

    "github.com/status-im/status-go/services/wallet/common"
)

type LatestBlockData struct {
    blockNumber   uint64
    timestamp     time.Time
    blockDuration time.Duration
}

type BlockChainState struct {
    blkMu              sync.RWMutex
    latestBlockNumbers map[uint64]LatestBlockData
    sinceFn            func(time.Time) time.Duration
}

func NewBlockChainState() *BlockChainState {
    return &BlockChainState{
        blkMu:              sync.RWMutex{},
        latestBlockNumbers: make(map[uint64]LatestBlockData),
        sinceFn:            time.Since,
    }
}

func (s *BlockChainState) GetEstimatedLatestBlockNumber(ctx context.Context, chainID uint64) (uint64, error) {
    blockNumber, _ := s.estimateLatestBlockNumber(chainID)
    return blockNumber, nil
}

func (s *BlockChainState) SetLastBlockNumber(chainID uint64, blockNumber uint64) {
    blockDuration, found := common.AverageBlockDurationForChain[common.ChainID(chainID)]
    if !found {
        blockDuration = common.AverageBlockDurationForChain[common.ChainID(common.UnknownChainID)]
    }
    s.setLatestBlockDataForChain(chainID, LatestBlockData{
        blockNumber:   blockNumber,
        timestamp:     time.Now(),
        blockDuration: blockDuration,
    })
}

func (s *BlockChainState) setLatestBlockDataForChain(chainID uint64, latestBlockData LatestBlockData) {
    s.blkMu.Lock()
    defer s.blkMu.Unlock()
    s.latestBlockNumbers[chainID] = latestBlockData
}

func (s *BlockChainState) estimateLatestBlockNumber(chainID uint64) (uint64, bool) {
    s.blkMu.RLock()
    defer s.blkMu.RUnlock()
    blockData, ok := s.latestBlockNumbers[chainID]
    if !ok {
        return 0, false
    }
    timeDiff := s.sinceFn(blockData.timestamp)
    return blockData.blockNumber + uint64((timeDiff / blockData.blockDuration)), true
}