Add lock to stream pull to prevent multiple nodes starting pulling for the same stream (#1187)

Author: leszko

go.mod

@@ -17,7 +17,7 @@ require (
 	github.com/hashicorp/serf v0.10.1
 	github.com/julienschmidt/httprouter v1.3.0
 	github.com/lib/pq v1.10.9
-	github.com/livepeer/go-api-client v0.4.18
+	github.com/livepeer/go-api-client v0.4.19
 	github.com/livepeer/go-tools v0.3.6
 	github.com/livepeer/joy4 v0.1.1
 	github.com/livepeer/livepeer-data v0.8.1

go.sum

@@ -451,6 +451,12 @@ github.com/livepeer/go-api-client v0.4.18-0.20240305122931-8f6d8c6543ad h1:eSqYA
 github.com/livepeer/go-api-client v0.4.18-0.20240305122931-8f6d8c6543ad/go.mod h1:Jdb+RI7JyzEZOHd1GUuKofwFDKMO/btTa80SdpUpYQw=
 github.com/livepeer/go-api-client v0.4.18 h1:grDZK6oMBm/6N9ZqsAk4ae2ohGDVzRPd2U12DrTRv2s=
 github.com/livepeer/go-api-client v0.4.18/go.mod h1:Jdb+RI7JyzEZOHd1GUuKofwFDKMO/btTa80SdpUpYQw=
+github.com/livepeer/go-api-client v0.4.19-0.20240321152123-03b5d097f6f0 h1:qJAVOO2lzdRSuL4GBgvpvSCXA/OUq1yOxpwZ9Nv/Xis=
+github.com/livepeer/go-api-client v0.4.19-0.20240321152123-03b5d097f6f0/go.mod h1:Jdb+RI7JyzEZOHd1GUuKofwFDKMO/btTa80SdpUpYQw=
+github.com/livepeer/go-api-client v0.4.19-0.20240325114751-e7cb002ff24d h1:DeO2VY8L2trThto7tc5KyiSjq+rTyLTnEz8je73kv6E=
+github.com/livepeer/go-api-client v0.4.19-0.20240325114751-e7cb002ff24d/go.mod h1:Jdb+RI7JyzEZOHd1GUuKofwFDKMO/btTa80SdpUpYQw=
+github.com/livepeer/go-api-client v0.4.19 h1:9YBSdYlYhCZdap08mIvptOR9B3Gf46AQOrWObJhvlkA=
+github.com/livepeer/go-api-client v0.4.19/go.mod h1:Jdb+RI7JyzEZOHd1GUuKofwFDKMO/btTa80SdpUpYQw=
 github.com/livepeer/go-tools v0.3.6 h1:LhRnoVVGFCtfBh6WyKdwJ2bPD/h5gaRvsAszmCqKt1Q=
 github.com/livepeer/go-tools v0.3.6/go.mod h1:qs31y68b3PQPmSr8nR8l5WQiIWI623z6pqOccqebjos=
 github.com/livepeer/joy4 v0.1.1 h1:Tz7gVcmvpG/nfUKHU+XJn6Qke/k32mTWMiH9qB0bhnM=

handlers/geolocation/geolocation.go

@@ -2,13 +2,15 @@ package geolocation
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"math/rand"
 	"net/http"
 	"net/url"
 	"regexp"
 	"strconv"
 	"strings"
+	"time"
 
 	"github.com/golang/glog"
 	"github.com/julienschmidt/httprouter"
@@ -20,6 +22,14 @@ import (
 	"github.com/livepeer/go-api-client"
 )
 
+const (
+	streamSourceRetries       = 20
+	streamSourceRetryInterval = 3 * time.Second
+	lockPullLeaseTimeout      = 3 * time.Minute
+)
+
+var errLockPull = errors.New("failed to lock pull")
+
 type GeolocationHandlersCollection struct {
 	Balancer balancer.Balancer
 	Cluster  cluster.Cluster
@@ -150,31 +160,52 @@ func (c *GeolocationHandlersCollection) HandleStreamSource(ctx context.Context,
 
 	latStr := fmt.Sprintf("%f", lat)
 	lonStr := fmt.Sprintf("%f", lon)
-	dtscURL, err := c.Balancer.MistUtilLoadSource(context.Background(), payload.StreamName, latStr, lonStr)
-	if err != nil {
-		glog.Errorf("error querying mist for STREAM_SOURCE: %s", err)
+	for i := 0; i < streamSourceRetries; i++ {
+		dtscURL, err := c.Balancer.MistUtilLoadSource(context.Background(), payload.StreamName, latStr, lonStr)
+		if err == nil {
+			return c.resolveReplicatedStream(dtscURL, payload.StreamName)
+		}
 
-		playbackID := payload.StreamName
-		parts := strings.Split(playbackID, "+")
-		if len(parts) == 2 {
-			playbackID = parts[1] // take the playbackID after the prefix e.g. 'video+'
+		glog.Errorf("error querying mist for STREAM_SOURCE: %s", err)
+		pullURL, err := c.getStreamPull(playbackIdFor(payload.StreamName))
+		if err == nil {
+			if pullURL == "" {
+				// not a stream pull
+				return "push://", nil
+			} else {
+				// start stream pull
+				glog.Infof("replying to Mist STREAM_SOURCE with stream pull request=%s response=%s", payload.StreamName, pullURL)
+				return pullURL, nil
+			}
 		}
-		pullURL, err := c.getStreamPull(playbackID)
-		if err != nil {
+		if !errors.Is(err, errLockPull) {
+			// stream pull failed for unknown reason
 			glog.Errorf("getStreamPull failed for %s: %s", payload.StreamName, err)
-		} else if pullURL != "" {
-			glog.Infof("replying to Mist STREAM_SOURCE with stream pull request=%s response=%s", payload.StreamName, pullURL)
-			return pullURL, nil
+			return "push://", nil
 		}
+		// stream pull failed, because another node started to pull at the same time
+		glog.Warningf("another node is currently pulling the stream, waiting %v and retrying", streamSourceRetryInterval)
+		time.Sleep(streamSourceRetryInterval)
+	}
+	return "push://", nil
+}
 
-		return "push://", nil
+func playbackIdFor(streamName string) string {
+	res := streamName
+	parts := strings.Split(res, "+")
+	if len(parts) == 2 {
+		res = parts[1] // take the playbackID after the prefix e.g. 'video+'
 	}
+	return res
+}
+
+func (c *GeolocationHandlersCollection) resolveReplicatedStream(dtscURL string, streamName string) (string, error) {
 	outURL, err := c.Cluster.ResolveNodeURL(dtscURL)
 	if err != nil {
 		glog.Errorf("error finding STREAM_SOURCE: %s", err)
 		return "push://", nil
 	}
-	glog.V(7).Infof("replying to Mist STREAM_SOURCE request=%s response=%s", payload.StreamName, outURL)
+	glog.V(7).Infof("replying to Mist STREAM_SOURCE request=%s response=%s", streamName, outURL)
 	return outURL, nil
 }
 
@@ -200,6 +231,10 @@ func (c *GeolocationHandlersCollection) getStreamPull(playbackID string) (string
 		return "", nil
 	}
 
+	if err := c.Lapi.LockPull(stream.ID, lockPullLeaseTimeout); err != nil {
+		return "", errLockPull
+	}
+
 	if len(stream.Pull.Headers) == 0 {
 		return stream.Pull.Source, nil
 	}