Merge pull request #6 from rusq/refactor

Refactor spaghetti file downloader.  Make it lasagna.

Author: rusq

files.go

@@ -41,7 +41,7 @@ func (sd *SlackDumper) filesFromMessages(m []Message) []slack.File {
 	return files
 }
 
-// SaveFileTo saves file to the specified directory
+// SaveFileTo saves file to the specified directory.
 func (sd *SlackDumper) SaveFileTo(dir string, f *slack.File) (int64, error) {
 	filePath := filepath.Join(dir, filename(f))
 	file, err := os.Create(filePath)
@@ -62,68 +62,78 @@ func filename(f *slack.File) string {
 	return fmt.Sprintf("%s-%s", f.ID, f.Name)
 }
 
-func (sd *SlackDumper) fileDownloader(dir string, files <-chan *slack.File, done chan<- bool) error {
-	const parallelDls = 4
+// fileDownloader will downloadstarts an sd.numDownloaders goroutines to
+// download files in parallel.  It will download any files that were received on toDownload channel,
+// and will close "done" once all downloads are complete.
+func (sd *SlackDumper) fileDownloader(dir string, toDownload <-chan *slack.File) (chan struct{}, error) {
+	done := make(chan struct{})
 
-	var wg sync.WaitGroup
-	dlQ := make(chan *slack.File)
-	stop := make(chan bool)
-
-	// downloaded contains file ids that already been downloaded,
-	// so we don't download the same file twice
-	downloaded := make(map[string]bool)
-
-	defer close(done)
+	if !sd.options.dumpfiles {
+		// terminating if dumpfiles is not enabled.
+		close(done)
+		return done, nil
+	}
 
 	if err := os.Mkdir(dir, 0777); err != nil {
 		if !os.IsExist(err) {
 			// channels done is closed by defer
-			return err
+			return done, err
 		}
 	}
 
-	worker := func(fs <-chan *slack.File) {
-	LOOP:
-		for {
-			select {
-			case file := <-fs:
-				// download file
-				log.Printf("saving %s, size: %d", filename(file), file.Size)
-				n, err := sd.SaveFileTo(dir, file)
-				if err != nil {
-					log.Printf("error saving %q: %s", filename(file), err)
-				}
-				log.Printf("file %s saved: %d bytes written", filename(file), n)
-			case <-stop:
-				break LOOP
-			}
+	var wg sync.WaitGroup
+	go func() {
+		// create workers
+		for i := 0; i < sd.options.workers; i++ {
+			wg.Add(1)
+			go func() {
+				sd.worker(dir, seenFilter(toDownload))
+				wg.Done()
+			}()
 		}
-		wg.Done()
-	}
+	}()
 
-	// create workers
-	for i := 0; i < parallelDls; i++ {
-		wg.Add(1)
-		go worker(dlQ)
-	}
+	// sentinel
+	go func() {
+		wg.Wait()
+		close(done)
+	}()
+
+	return done, nil
+}
 
-	// files queue must be closed on the sender side (see DumpToDir.(1))
-	for f := range files {
-		_, ok := downloaded[f.ID]
-		if ok {
-			log.Printf("already seen %s, skipping", filename(f))
-			continue
+func (sd *SlackDumper) worker(dir string, filesC <-chan *slack.File) {
+	for file := range filesC {
+		// download file
+		log.Printf("saving %s, size: %d", filename(file), file.Size)
+		n, err := sd.SaveFileTo(dir, file)
+		if err != nil {
+			log.Printf("error saving %q: %s", filename(file), err)
 		}
-		dlQ <- f
-		downloaded[f.ID] = true
+		log.Printf("file %s saved: %d bytes written", filename(file), n)
 	}
+}
 
-	// closing stop will terminate all workers (1)
-	close(stop)
-	// workers mark all WorkGroups as done (2)
-	wg.Wait()
-	// we send the signal to caller that we're done too
-	done <- true
-
-	return nil
+// seenFilter filters the files from filesC to ensure that no duplicates
+// are downloaded.
+func seenFilter(filesC <-chan *slack.File) <-chan *slack.File {
+	dlQ := make(chan *slack.File)
+	go func() {
+		// closing stop will lead to all worker goroutines to terminate.
+		defer close(dlQ)
+
+		// seen contains file ids that already been seen,
+		// so we don't download the same file twice
+		seen := make(map[string]bool)
+		// files queue must be closed by the caller (see DumpToDir.(1))
+		for f := range filesC {
+			if _, ok := seen[f.ID]; ok {
+				log.Printf("already seen %s, skipping", filename(f))
+				continue
+			}
+			seen[f.ID] = true
+			dlQ <- f
+		}
+	}()
+	return dlQ
 }

go.mod

@@ -1,13 +1,20 @@
 module github.com/rusq/slackdump
 
-go 1.13
+go 1.17
 
 require (
 	github.com/joho/godotenv v1.4.0
 	github.com/pkg/errors v0.9.1
 	github.com/slack-go/slack v0.9.5
-	github.com/stretchr/testify v1.4.0 // indirect
+	github.com/stretchr/testify v1.4.0
 	golang.org/x/time v0.0.0-20211116232009-f0f3c7e86c11
 )
 
+require (
+	github.com/davecgh/go-spew v1.1.1 // indirect
+	github.com/gorilla/websocket v1.4.2 // indirect
+	github.com/pmezard/go-difflib v1.0.0 // indirect
+	gopkg.in/yaml.v2 v2.2.2 // indirect
+)
+
 replace github.com/slack-go/slack => github.com/rusq/slack v0.9.6

slackdump.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"io"
 	"log"
+	"runtime"
 	"sort"
 
 	"github.com/slack-go/slack"
@@ -25,6 +26,7 @@ type SlackDumper struct {
 
 type options struct {
 	dumpfiles bool
+	workers   int
 }
 
 var allChanTypes = []string{"mpim", "im", "public_channel", "private_channel"}
@@ -42,12 +44,28 @@ func DumpFiles(b bool) Option {
 	}
 }
 
+const defNumWorkers = 4 // seems reasonable
+
+// NumWorkers allows to set the number of file download workers. n should be in
+// range [1, NumCPU]. If not in range, will be reset to a defNumWorkers number,
+// which seems reasonable.
+func NumWorkers(n int) Option {
+	return func(sd *SlackDumper) {
+		if n < 1 || runtime.NumCPU() < n {
+			n = defNumWorkers
+		}
+		sd.options.workers = n
+	}
+}
+
 // New creates new client and populates the internal cache of users and channels
 // for lookups.
 func New(ctx context.Context, token string, cookie string, opts ...Option) (*SlackDumper, error) {
-	var err error
 	sd := &SlackDumper{
 		client: slack.New(token, slack.OptionCookie(cookie)),
+		options: options{
+			workers: defNumWorkers,
+		},
 	}
 	for _, opt := range opts {
 		opt(sd)
@@ -59,6 +77,7 @@ func New(ctx context.Context, token string, cookie string, opts ...Option) (*Sla
 
 	go func() {
 		defer close(errC)
+
 		var err error
 		chanTypes := allChanTypes
 		log.Println("> caching channels, might take a while...")
@@ -73,7 +92,7 @@ func New(ctx context.Context, token string, cookie string, opts ...Option) (*Sla
 		return nil, fmt.Errorf("error fetching users: %s", err)
 	}
 
-	if err = <-errC; err != nil {
+	if err := <-errC; err != nil {
 		return nil, fmt.Errorf("error fetching channels: %s", err)
 	}
 
@@ -93,43 +112,36 @@ func (sd *SlackDumper) IsDeletedUser(id string) bool {
 
 // DumpMessages fetches messages from the conversation identified by channelID.
 func (sd *SlackDumper) DumpMessages(ctx context.Context, channelID string) (*Channel, error) {
+	var filesC = make(chan *slack.File, 20)
 
-	params := &slack.GetConversationHistoryParameters{
-		ChannelID: channelID,
-	}
-
-	filesC := make(chan *slack.File, 20)
-	done := make(chan bool)
-	errC := make(chan error, 1)
-
-	if sd.options.dumpfiles {
-		go func() {
-			errC <- sd.fileDownloader(channelID, filesC, done)
-		}()
+	dlDoneC, err := sd.fileDownloader(channelID, filesC)
+	if err != nil {
+		return nil, err
 	}
 
 	limiter := newLimiter(tier3)
-	var messages []Message
 
+	var (
+		messages []Message
+		cursor   string
+	)
 	for i := 1; ; i++ {
-		select {
-		case err := <-errC:
-			// stop the goroutine gracefully if it's running
-			close(filesC)
-			<-done
-			return nil, err
-		default:
-		}
-		resp, err := sd.client.GetConversationHistoryContext(ctx, params)
+		resp, err := sd.client.GetConversationHistoryContext(
+			ctx,
+			&slack.GetConversationHistoryParameters{
+				ChannelID: channelID,
+				Cursor:    cursor,
+			},
+		)
 		if err != nil {
 			return nil, err
 		}
-		chunk := sd.convertMsgs(resp.Messages)
 
+		chunk := sd.convertMsgs(resp.Messages)
 		if err := sd.populateThreads(ctx, chunk, channelID, limiter); err != nil {
 			return nil, err
 		}
-		sd.extractFiles(filesC, chunk)
+		sd.pipeFiles(filesC, chunk)
 		messages = append(messages, chunk...)
 
 		log.Printf("request #%d, fetched: %d, total: %d\n",
@@ -139,7 +151,7 @@ func (sd *SlackDumper) DumpMessages(ctx context.Context, channelID string) (*Cha
 			break
 		}
 
-		params.Cursor = resp.ResponseMetaData.NextCursor
+		cursor = resp.ResponseMetaData.NextCursor
 
 		limiter.Wait(ctx)
 	}
@@ -150,7 +162,7 @@ func (sd *SlackDumper) DumpMessages(ctx context.Context, channelID string) (*Cha
 
 	if sd.options.dumpfiles {
 		close(filesC)
-		<-done
+		<-dlDoneC
 	}
 
 	return &Channel{Messages: messages, ID: channelID}, nil
@@ -165,21 +177,22 @@ func (sd *SlackDumper) convertMsgs(sm []slack.Message) []Message {
 	return msgs
 }
 
-// extractFiles scans the messages and sends all the files discovered to the filesC.
-func (sd *SlackDumper) extractFiles(filesC chan<- *slack.File, msgs []Message) {
+// pipeFiles scans the messages and sends all the files discovered to the filesC.
+func (sd *SlackDumper) pipeFiles(filesC chan<- *slack.File, msgs []Message) {
 	if !sd.options.dumpfiles {
 		return
 	}
 	// place files in download queue
-	chunk := sd.filesFromMessages(msgs)
-	for i := range chunk {
-		filesC <- &chunk[i]
+	fileChunk := sd.filesFromMessages(msgs)
+	for i := range fileChunk {
+		filesC <- &fileChunk[i]
 	}
 }
 
 // populateThreads scans the message slice for threads, if and when it
 // discovers the message with ThreadTimestamp, it fetches all messages in that
 // thread updating them to the msgs slice.
+//
 // ref: https://api.slack.com/messaging/retrieving
 func (sd *SlackDumper) populateThreads(ctx context.Context, msgs []Message, channelID string, l *rate.Limiter) error {
 	for i := range msgs {