downloader: added downloaded torrents notifier

erigon-lib/direct/downloader_client.go

@@ -18,6 +18,7 @@ package direct
 
 import (
 	"context"
+	"io"
 
 	proto_downloader "github.com/erigontech/erigon-lib/gointerfaces/downloaderproto"
 	"google.golang.org/grpc"
@@ -51,3 +52,65 @@ func (c *DownloaderClient) SetLogPrefix(ctx context.Context, in *proto_downloade
 func (c *DownloaderClient) Completed(ctx context.Context, in *proto_downloader.CompletedRequest, opts ...grpc.CallOption) (*proto_downloader.CompletedReply, error) {
 	return c.server.Completed(ctx, in)
 }
+
+func (c *DownloaderClient) TorrentCompleted(ctx context.Context, in *proto_downloader.TorrentCompletedRequest, opts ...grpc.CallOption) (proto_downloader.Downloader_TorrentCompletedClient, error) {
+	ch := make(chan *downloadedReply, 1<<16)
+	streamServer := &DownloadeSubscribeS{ch: ch, ctx: ctx}
+
+	go func() {
+		streamServer.Err(c.server.TorrentCompleted(in, streamServer))
+	}()
+
+	return &DownloadeSubscribeC{ch: ch, ctx: ctx}, nil
+}
+
+type DownloadeSubscribeC struct {
+	ch  chan *downloadedReply
+	ctx context.Context
+	grpc.ClientStream
+}
+
+func (c *DownloadeSubscribeC) Recv() (*proto_downloader.TorrentCompletedReply, error) {
+	if c.ctx.Err() != nil {
+		return nil, io.EOF
+	}
+
+	m, ok := <-c.ch
+	if !ok || m == nil {
+		return nil, io.EOF
+	}
+	return m.r, m.err
+}
+func (c *DownloadeSubscribeC) Context() context.Context { return c.ctx }
+
+type DownloadeSubscribeS struct {
+	ch  chan *downloadedReply
+	ctx context.Context
+	grpc.ServerStream
+}
+
+type downloadedReply struct {
+	r   *proto_downloader.TorrentCompletedReply
+	err error
+}
+
+func (s *DownloadeSubscribeS) Send(m *proto_downloader.TorrentCompletedReply) error {
+	if s.ctx.Err() != nil {
+		if s.ch != nil {
+			ch := s.ch
+			s.ch = nil
+			close(ch)
+		}
+		return s.ctx.Err()
+	}
+
+	s.ch <- &downloadedReply{r: m}
+	return nil
+}
+func (s *DownloadeSubscribeS) Context() context.Context { return s.ctx }
+func (s *DownloadeSubscribeS) Err(err error) {
+	if err == nil {
+		return
+	}
+	s.ch <- &downloadedReply{err: err}
+}

erigon-lib/downloader/downloader.go

@@ -59,6 +59,7 @@ import (
 	"github.com/erigontech/erigon-lib/diagnostics"
 	"github.com/erigontech/erigon-lib/downloader/downloadercfg"
 	"github.com/erigontech/erigon-lib/downloader/snaptype"
+	prototypes "github.com/erigontech/erigon-lib/gointerfaces/typesproto"
 	"github.com/erigontech/erigon-lib/kv"
 	"github.com/erigontech/erigon-lib/kv/mdbx"
 	"github.com/erigontech/erigon-lib/log/v3"
@@ -97,8 +98,15 @@ type Downloader struct {
 
 	stuckFileDetailedLogs bool
 
-	logPrefix string
-	startTime time.Time
+	logPrefix         string
+	startTime         time.Time
+	broadcast         func(name string, hash *prototypes.H160)
+	completedTorrents map[string]completedTorrentInfo
+}
+
+type completedTorrentInfo struct {
+	path string
+	hash *prototypes.H160
 }
 
 type downloadInfo struct {
@@ -345,6 +353,7 @@ func New(ctx context.Context, cfg *downloadercfg.Cfg, logger log.Logger, verbosi
 		downloading:         map[string]*downloadInfo{},
 		webseedsDiscover:    discover,
 		logPrefix:           "",
+		completedTorrents:   make(map[string]completedTorrentInfo),
 	}
 	d.webseeds.SetTorrent(d.torrentFS, snapLock.Downloads, cfg.DownloadTorrentFilesFromWebseed)
 
@@ -906,6 +915,7 @@ func (d *Downloader) mainLoop(silent bool) error {
 			for _, t := range torrents {
 				if _, ok := complete[t.Name()]; ok {
 					clist = append(clist, t.Name())
+					d.torrentCompleted(t.Name(), t.InfoHash())
 					continue
 				}
 
@@ -959,6 +969,7 @@ func (d *Downloader) mainLoop(silent bool) error {
 							} else {
 								clist = append(clist, t.Name())
 								complete[t.Name()] = struct{}{}
+								d.torrentCompleted(t.Name(), t.InfoHash())
 								continue
 							}
 						}
@@ -1025,6 +1036,8 @@ func (d *Downloader) mainLoop(silent bool) error {
 					d.lock.Unlock()
 					complete[t.Name()] = struct{}{}
 					clist = append(clist, t.Name())
+					d.torrentCompleted(t.Name(), t.InfoHash())
+
 					continue
 				}
 
@@ -2762,6 +2775,7 @@ func (d *Downloader) BuildTorrentFilesIfNeed(ctx context.Context, chain string,
 	_, err := BuildTorrentFilesIfNeed(ctx, d.cfg.Dirs, d.torrentFS, chain, ignore, false)
 	return err
 }
+
 func (d *Downloader) Stats() AggStats {
 	d.lock.RLock()
 	defer d.lock.RUnlock()
@@ -2946,3 +2960,32 @@ func calculateTime(amountLeft, rate uint64) string {
 func (d *Downloader) Completed() bool {
 	return d.stats.Completed
 }
+
+// Store completed torrents in order to notify GrpcServer subscribers when they subscribe and there is already downloaded files
+func (d *Downloader) torrentCompleted(tName string, tHash metainfo.Hash) {
+	d.lock.Lock()
+	defer d.lock.Unlock()
+	hash := InfoHashes2Proto(tHash)
+
+	//check is torrent already completed cause some funcs may call this method multiple times
+	if _, ok := d.completedTorrents[tName]; !ok {
+		d.notifyCompleted(tName, hash)
+	}
+
+	d.completedTorrents[tName] = completedTorrentInfo{
+		path: tName,
+		hash: hash,
+	}
+}
+
+// Notify GrpcServer subscribers about completed torrent
+func (d *Downloader) notifyCompleted(tName string, tHash *prototypes.H160) {
+	d.broadcast(tName, tHash)
+}
+
+func (d *Downloader) getCompletedTorrents() map[string]completedTorrentInfo {
+	d.lock.RLock()
+	defer d.lock.RUnlock()
+
+	return d.completedTorrents
+}

erigon-lib/downloader/downloader_grpc_server.go

@@ -22,6 +22,8 @@ import (
 	"fmt"
 	"os"
 	"path/filepath"
+	"slices"
+	"sync"
 	"time"
 
 	"github.com/anacrolix/torrent/metainfo"
@@ -38,12 +40,20 @@ var (
 )
 
 func NewGrpcServer(d *Downloader) (*GrpcServer, error) {
-	return &GrpcServer{d: d}, nil
+	svr := &GrpcServer{
+		d: d,
+	}
+
+	d.broadcast = svr.broadcast
+
+	return svr, nil
 }
 
 type GrpcServer struct {
 	proto_downloader.UnimplementedDownloaderServer
-	d *Downloader
+	d           *Downloader
+	mu          sync.RWMutex
+	subscribers []proto_downloader.Downloader_TorrentCompletedServer
 }
 
 func (s *GrpcServer) ProhibitNewDownloads(ctx context.Context, req *proto_downloader.ProhibitNewDownloadsRequest) (*emptypb.Empty, error) {
@@ -129,6 +139,10 @@ func Proto2InfoHash(in *prototypes.H160) metainfo.Hash {
 	return gointerfaces.ConvertH160toAddress(in)
 }
 
+func InfoHashes2Proto(in metainfo.Hash) *prototypes.H160 {
+	return gointerfaces.ConvertAddressToH160(in)
+}
+
 func (s *GrpcServer) SetLogPrefix(ctx context.Context, request *proto_downloader.SetLogPrefixRequest) (*emptypb.Empty, error) {
 	s.d.SetLogPrefix(request.Prefix)
 
@@ -138,3 +152,41 @@ func (s *GrpcServer) SetLogPrefix(ctx context.Context, request *proto_downloader
 func (s *GrpcServer) Completed(ctx context.Context, request *proto_downloader.CompletedRequest) (*proto_downloader.CompletedReply, error) {
 	return &proto_downloader.CompletedReply{Completed: s.d.Completed()}, nil
 }
+
+func (s *GrpcServer) TorrentCompleted(req *proto_downloader.TorrentCompletedRequest, stream proto_downloader.Downloader_TorrentCompletedServer) error {
+	// Register the new subscriber
+	s.mu.Lock()
+	s.subscribers = append(s.subscribers, stream)
+	s.mu.Unlock()
+
+	//Notifying about all completed torrents to the new subscriber
+	cmp := s.d.getCompletedTorrents()
+	for _, cmpInfo := range cmp {
+		s.broadcast(cmpInfo.path, cmpInfo.hash)
+	}
+
+	return nil
+}
+
+func (s *GrpcServer) broadcast(name string, hash *prototypes.H160) {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+
+	var unsub []int
+
+	for i, s := range s.subscribers {
+		if s.Context().Err() != nil {
+			unsub = append(unsub, i)
+			continue
+		}
+
+		s.Send(&proto_downloader.TorrentCompletedReply{
+			Name: name,
+			Hash: hash,
+		})
+	}
+
+	for i := len(unsub) - 1; i >= 0; i-- {
+		s.subscribers = slices.Delete(s.subscribers, unsub[i], unsub[i])
+	}
+}

erigon-lib/downloader/downloadergrpc/client.go

@@ -80,3 +80,8 @@ func String2Proto(in string) *prototypes.H160 {
 	copy(infoHash[:], inHex)
 	return gointerfaces.ConvertAddressToH160(infoHash)
 }
+
+func Proto2String(in *prototypes.H160) string {
+	addr := gointerfaces.ConvertH160toAddress(in)
+	return hex.EncodeToString(addr[:])
+}