waku-org/go-waku

package history

import (
    "context"
    "crypto/rand"
    "errors"
    "fmt"
    "math"
    "math/big"
    "net"
    "net/http"
    "runtime"
    "sort"
    "sync"
    "time"

    "github.com/libp2p/go-libp2p/core/peer"
    "github.com/waku-org/go-waku/waku/v2/api/common"
    "github.com/waku-org/go-waku/waku/v2/protocol/store"
    "go.uber.org/zap"
)

const defaultBackoff = 10 * time.Second
const graylistBackoff = 3 * time.Minute
const storenodeVerificationInterval = time.Second
const storenodeMaxFailedRequests uint = 2
const minStorenodesToChooseFrom = 3
const isAndroidEmulator = runtime.GOOS == "android" && runtime.GOARCH == "amd64"
const findNearestMailServer = !isAndroidEmulator
const overrideDNS = runtime.GOOS == "android" || runtime.GOOS == "ios"
const bootstrapDNS = "8.8.8.8:53"

type connStatus int

const (
    disconnected connStatus = iota + 1
    connected
)

type peerStatus struct {
    status                connStatus
    canConnectAfter       time.Time
    lastConnectionAttempt time.Time
}

type StorenodeConfigProvider interface {
    UseStorenodes() (bool, error)
    GetPinnedStorenode() (peer.AddrInfo, error)
    Storenodes() ([]peer.AddrInfo, error)
}

type StorenodeCycle struct {
    sync.RWMutex

    logger *zap.Logger

    storenodeConfigProvider StorenodeConfigProvider
    pinger                  common.Pinger

    StorenodeAvailableOneshotEmitter *OneShotEmitter[struct{}]
    StorenodeChangedEmitter          *Emitter[peer.ID]
    StorenodeNotWorkingEmitter       *Emitter[struct{}]
    StorenodeAvailableEmitter        *Emitter[peer.ID]

    failedRequests map[peer.ID]uint

    peersMutex      sync.RWMutex
    activeStorenode peer.ID
    peers           map[peer.ID]peerStatus
}

func NewStorenodeCycle(logger *zap.Logger, pinger common.Pinger) *StorenodeCycle {
    return &StorenodeCycle{
        StorenodeAvailableOneshotEmitter: NewOneshotEmitter[struct{}](),
        StorenodeChangedEmitter:          NewEmitter[peer.ID](),
        StorenodeNotWorkingEmitter:       NewEmitter[struct{}](),
        StorenodeAvailableEmitter:        NewEmitter[peer.ID](),
        logger:                           logger.Named("storenode-cycle"),
    }
}

func (m *StorenodeCycle) Start(ctx context.Context) {
    m.logger.Debug("starting storenode cycle")
    m.failedRequests = make(map[peer.ID]uint)
    m.peers = make(map[peer.ID]peerStatus)

    go m.verifyStorenodeStatus(ctx)
}

func (m *StorenodeCycle) DisconnectActiveStorenode(backoff time.Duration) {
    m.Lock()
    defer m.Unlock()

    m.disconnectActiveStorenode(backoff)
}

func (m *StorenodeCycle) connectToNewStorenodeAndWait(ctx context.Context) error {
    // Handle pinned storenodes
    m.logger.Info("disconnecting storenode")
    pinnedStorenode, err := m.storenodeConfigProvider.GetPinnedStorenode()
    if err != nil {
        m.logger.Error("could not obtain the pinned storenode", zap.Error(err))
        return err
    }

    // If no pinned storenode, no need to disconnect and wait for it to be available
    if pinnedStorenode.ID == "" {
        m.disconnectActiveStorenode(graylistBackoff)
    }

    return m.findNewStorenode(ctx)
}

func (m *StorenodeCycle) disconnectStorenode(backoffDuration time.Duration) error {
    if m.activeStorenode == "" {
        m.logger.Info("no active storenode")
        return nil
    }

    m.logger.Info("disconnecting active storenode", zap.Stringer("peerID", m.activeStorenode))

    m.peersMutex.Lock()
    pInfo, ok := m.peers[m.activeStorenode]
    if ok {
        pInfo.status = disconnected
        pInfo.canConnectAfter = time.Now().Add(backoffDuration)
        m.peers[m.activeStorenode] = pInfo
    } else {
        m.peers[m.activeStorenode] = peerStatus{
            status:          disconnected,
            canConnectAfter: time.Now().Add(backoffDuration),
        }
    }
    m.peersMutex.Unlock()

    m.activeStorenode = ""

    return nil
}

func (m *StorenodeCycle) disconnectActiveStorenode(backoffDuration time.Duration) {
    err := m.disconnectStorenode(backoffDuration)
    if err != nil {
        m.logger.Error("failed to disconnect storenode", zap.Error(err))
    }

    m.StorenodeChangedEmitter.Emit("")
}

func (m *StorenodeCycle) Cycle(ctx context.Context) {
    if m.storenodeConfigProvider == nil {
        m.logger.Debug("storenodeConfigProvider not yet setup")
        return
    }

    m.logger.Info("Automatically switching storenode")

    if m.activeStorenode != "" {
        m.disconnectActiveStorenode(graylistBackoff)
    }

    useStorenode, err := m.storenodeConfigProvider.UseStorenodes()
    if err != nil {
        m.logger.Error("failed to get use storenodes", zap.Error(err))
        return
    }

    if !useStorenode {
        m.logger.Info("Skipping storenode search due to useStorenode being false")
        return
    }

    err = m.findNewStorenode(ctx)
    if err != nil {
        m.logger.Error("Error getting new storenode", zap.Error(err))
    }
}

func poolSize(fleetSize int) int {
    return int(math.Ceil(float64(fleetSize) / 4))
}

func (m *StorenodeCycle) getAvailableStorenodesSortedByRTT(ctx context.Context, allStorenodes []peer.AddrInfo) []peer.AddrInfo {
    peerIDToInfo := make(map[peer.ID]peer.AddrInfo)
    for _, p := range allStorenodes {
        peerIDToInfo[p.ID] = p
    }

    availableStorenodes := make(map[peer.ID]time.Duration)
    availableStorenodesMutex := sync.Mutex{}
    availableStorenodesWg := sync.WaitGroup{}
    for _, storenode := range allStorenodes {
        availableStorenodesWg.Add(1)
        go func(peerInfo peer.AddrInfo) {
            defer availableStorenodesWg.Done()
            ctx, cancel := context.WithTimeout(ctx, 4*time.Second)
            defer cancel()

            rtt, err := m.pinger.PingPeer(ctx, peerInfo)
            if err == nil { // pinging storenodes might fail, but we don't care
                availableStorenodesMutex.Lock()
                availableStorenodes[peerInfo.ID] = rtt
                availableStorenodesMutex.Unlock()
            }
        }(storenode)
    }
    availableStorenodesWg.Wait()

    if len(availableStorenodes) == 0 {
        m.logger.Warn("No storenodes available") // Do nothing..
        return nil
    }

    var sortedStorenodes []SortedStorenode
    for storenodeID, rtt := range availableStorenodes {
        sortedStorenode := SortedStorenode{
            Storenode: peerIDToInfo[storenodeID],
            RTT:       rtt,
        }
        m.peersMutex.Lock()
        pInfo, ok := m.peers[storenodeID]
        m.peersMutex.Unlock()
        if ok && time.Now().Before(pInfo.canConnectAfter) {
            continue // We can't connect to this node yet
        }
        sortedStorenodes = append(sortedStorenodes, sortedStorenode)
    }
    sort.Sort(byRTTMsAndCanConnectBefore(sortedStorenodes))

    result := make([]peer.AddrInfo, len(sortedStorenodes))
    for i, s := range sortedStorenodes {
        result[i] = s.Storenode
    }

    return result
}

func (m *StorenodeCycle) findNewStorenode(ctx context.Context) error {
    // we have to override DNS manually because of https://github.com/status-im/status-mobile/issues/19581
    if overrideDNS {
        var dialer net.Dialer
        net.DefaultResolver = &net.Resolver{
            PreferGo: false,
            Dial: func(context context.Context, _, _ string) (net.Conn, error) {
                conn, err := dialer.DialContext(context, "udp", bootstrapDNS)
                if err != nil {
                    return nil, err
                }
                return conn, nil
            },
        }
    }

    pinnedStorenode, err := m.storenodeConfigProvider.GetPinnedStorenode()
    if err != nil {
        m.logger.Error("Could not obtain the pinned storenode", zap.Error(err))
        return err
    }

    if pinnedStorenode.ID != "" {
        return m.setActiveStorenode(pinnedStorenode.ID)
    }

    m.logger.Info("Finding a new storenode..")

    allStorenodes, err := m.storenodeConfigProvider.Storenodes()
    if err != nil {
        return err
    }

    //    TODO: remove this check once sockets are stable on x86_64 emulators
    if findNearestMailServer {
        allStorenodes = m.getAvailableStorenodesSortedByRTT(ctx, allStorenodes)
    }

    // Picks a random storenode amongs the ones with the lowest latency
    // The pool size is 1/4 of the storenodes were pinged successfully
    // If the pool size is less than `minStorenodesToChooseFrom`, it will
    // pick a storenode fromm all the available storenodes
    pSize := poolSize(len(allStorenodes) - 1)
    if pSize <= minStorenodesToChooseFrom {
        pSize = len(allStorenodes)
        if pSize <= 0 {
            m.logger.Warn("No storenodes available") // Do nothing..
            return nil
        }
    }

    r, err := rand.Int(rand.Reader, big.NewInt(int64(pSize)))
    if err != nil {
        return err
    }

    ms := allStorenodes[r.Int64()]
    return m.setActiveStorenode(ms.ID)
}

func (m *StorenodeCycle) storenodeStatus(peerID peer.ID) connStatus {
    m.peersMutex.RLock()
    defer m.peersMutex.RUnlock()

    peer, ok := m.peers[peerID]
    if !ok {
        return disconnected
    }
    return peer.status
}

func (m *StorenodeCycle) setActiveStorenode(peerID peer.ID) error {
    m.activeStorenode = peerID

    m.StorenodeChangedEmitter.Emit(m.activeStorenode)

    storenodeStatus := m.storenodeStatus(peerID)
    if storenodeStatus != connected {
        m.peersMutex.Lock()
        m.peers[peerID] = peerStatus{
            status:                connected,
            lastConnectionAttempt: time.Now(),
            canConnectAfter:       time.Now().Add(defaultBackoff),
        }
        m.peersMutex.Unlock()

        m.failedRequests[peerID] = 0
        m.logger.Info("storenode available", zap.Stringer("peerID", m.activeStorenode))

        m.StorenodeAvailableOneshotEmitter.Emit(struct{}{}) // Maybe can be refactored away?
        m.StorenodeAvailableEmitter.Emit(m.activeStorenode)
    }
    return nil
}

func (m *StorenodeCycle) GetActiveStorenode() peer.ID {
    m.RLock()
    defer m.RUnlock()

    return m.activeStorenode
}

func (m *StorenodeCycle) IsStorenodeAvailable(peerID peer.ID) bool {
    return m.storenodeStatus(peerID) == connected
}

func (m *StorenodeCycle) penalizeStorenode(id peer.ID) {
    m.peersMutex.Lock()
    defer m.peersMutex.Unlock()
    pInfo, ok := m.peers[id]
    if !ok {
        pInfo.status = disconnected
    }

    pInfo.canConnectAfter = time.Now().Add(graylistBackoff)
    m.peers[id] = pInfo
}

func (m *StorenodeCycle) verifyStorenodeStatus(ctx context.Context) {
    ticker := time.NewTicker(storenodeVerificationInterval)
    defer ticker.Stop()

    for {
        select {
        case <-ticker.C:
            err := m.disconnectStorenodeIfRequired(ctx)
            if err != nil {
                m.logger.Error("failed to handle storenode cycle event", zap.Error(err))
                continue
            }

        case <-ctx.Done():
            return
        }
    }
}

func (m *StorenodeCycle) disconnectStorenodeIfRequired(ctx context.Context) error {
    m.logger.Debug("wakuV2 storenode status verification")

    if m.activeStorenode == "" {
        // No active storenode, find a new one
        m.Cycle(ctx)
        return nil
    }

    // Check whether we want to disconnect the active storenode
    if m.failedRequests[m.activeStorenode] >= storenodeMaxFailedRequests {
        m.penalizeStorenode(m.activeStorenode)
        m.StorenodeNotWorkingEmitter.Emit(struct{}{})

        m.logger.Info("too many failed requests", zap.Stringer("storenode", m.activeStorenode))
        m.failedRequests[m.activeStorenode] = 0
        return m.connectToNewStorenodeAndWait(ctx)
    }

    return nil
}

func (m *StorenodeCycle) SetStorenodeConfigProvider(provider StorenodeConfigProvider) {
    m.storenodeConfigProvider = provider
}

func (m *StorenodeCycle) WaitForAvailableStoreNode(ctx context.Context) bool {
    // Note: Add 1 second to timeout, because the storenode cycle has 1 second ticker, which doesn't tick on start.
    // This can be improved after merging https://github.com/status-im/status-go/pull/4380.
    // NOTE: https://stackoverflow.com/questions/32705582/how-to-get-time-tick-to-tick-immediately

    wg := sync.WaitGroup{}
    wg.Add(1)
    go func() {
        defer wg.Done()
        for !m.IsStorenodeAvailable(m.activeStorenode) {
            select {
            case <-m.StorenodeAvailableOneshotEmitter.Subscribe():
            case <-ctx.Done():
                if errors.Is(ctx.Err(), context.Canceled) {
                    return
                }

                // Wait for an additional second, but handle cancellation
                select {
                case <-time.After(1 * time.Second):
                case <-ctx.Done(): // context was cancelled
                }

                return

            }
        }
    }()

    select {
    case <-waitForWaitGroup(&wg):
    case <-ctx.Done():
        // Wait for an additional second, but handle cancellation
        select {
        case <-time.After(1 * time.Second):
        case <-ctx.Done(): // context was cancelled o
        }
    }

    return m.IsStorenodeAvailable(m.activeStorenode)
}

func waitForWaitGroup(wg *sync.WaitGroup) <-chan struct{} {
    ch := make(chan struct{})
    go func() {
        wg.Wait()
        close(ch)
    }()
    return ch
}

type storenodeTaskParameters struct {
    customPeerID peer.ID
}

type StorenodeTaskOption func(*storenodeTaskParameters)

func WithPeerID(peerID peer.ID) StorenodeTaskOption {
    return func(stp *storenodeTaskParameters) {
        stp.customPeerID = peerID
    }
}

func (m *StorenodeCycle) PerformStorenodeTask(fn func() error, options ...StorenodeTaskOption) error {
    params := storenodeTaskParameters{}
    for _, opt := range options {
        opt(&params)
    }

    peerID := params.customPeerID
    if peerID == "" {
        peerID = m.GetActiveStorenode()
    }

    if peerID == "" {
        return errors.New("storenode not available")
    }

    m.RLock()
    defer m.RUnlock()

    var tries uint = 0
    for tries < storenodeMaxFailedRequests {
        if params.customPeerID == "" && m.storenodeStatus(peerID) != connected {
            return errors.New("storenode not available")
        }
        m.logger.Info("trying performing history requests", zap.Uint("try", tries), zap.Stringer("peerID", peerID))

        // Peform request
        err := fn()
        if err == nil {
            // Reset failed requests
            m.logger.Debug("history request performed successfully", zap.Stringer("peerID", peerID))
            m.failedRequests[peerID] = 0
            return nil
        }

        m.logger.Error("failed to perform history request",
            zap.Stringer("peerID", peerID),
            zap.Uint("tries", tries),
            zap.Error(err),
        )

        tries++

        if storeErr, ok := err.(*store.StoreError); ok {
            if storeErr.Code == http.StatusTooManyRequests {
                m.disconnectActiveStorenode(defaultBackoff)
                return fmt.Errorf("ratelimited at storenode %s: %w", peerID, err)
            }
        }

        // Increment failed requests
        m.failedRequests[peerID]++

        // Change storenode
        if m.failedRequests[peerID] >= storenodeMaxFailedRequests {
            return errors.New("too many failed requests")
        }
        // Wait a couple of second not to spam
        time.Sleep(2 * time.Second)

    }
    return errors.New("failed to perform history request")
}