Nuke x/sync

1 files changed, 30 insertions, 7 deletions

diff --git a/main.go b/main.go
index 7edf574..e86bb36 100644
--- a/main.go
+++ b/main.go
@@ -29,7 +29,6 @@ import (
 	"time"
 
 	"github.com/mattn/go-sqlite3"
-	"golang.org/x/sync/semaphore"
 )
 
 var (
@@ -37,7 +36,7 @@ var (
 	galleryDirectory string  // gallery directory
 
 	// taskSemaphore limits parallel computations.
-	taskSemaphore *semaphore.Weighted
+	taskSemaphore semaphore
 )
 
 func hammingDistance(a, b int64) int {
@@ -89,6 +88,30 @@ func dbCollectStrings(query string) ([]string, error) {
 	return result, nil
 }
 
+// --- Semaphore ---------------------------------------------------------------
+
+type semaphore chan struct{}
+
+func newSemaphore(size int) semaphore { return make(chan struct{}, size) }
+func (s semaphore) release()          { <-s }
+
+func (s semaphore) acquire(ctx context.Context) error {
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case s <- struct{}{}:
+	}
+
+	// Give priority to context cancellation.
+	select {
+	case <-ctx.Done():
+		s.release()
+		return ctx.Err()
+	default:
+	}
+	return nil
+}
+
 // --- Progress bar ------------------------------------------------------------
 
 type progressBar struct {
@@ -740,13 +763,13 @@ func cmdImport(args []string) error {
 	ctx, cancel := context.WithCancelCause(context.Background())
 	wg := sync.WaitGroup{}
 	for _, path := range paths {
-		if taskSemaphore.Acquire(ctx, 1) != nil {
+		if taskSemaphore.acquire(ctx) != nil {
 			break
 		}
 
 		wg.Add(1)
 		go func(path string) {
-			defer taskSemaphore.Release(1)
+			defer taskSemaphore.release()
 			defer wg.Done()
 			if err := i.Import(path); err != nil {
 				cancel(err)
@@ -1071,13 +1094,13 @@ func cmdThumbnail(args []string) error {
 	ctx, cancel := context.WithCancelCause(context.Background())
 	wg := sync.WaitGroup{}
 	for _, sha1 := range hexSHA1 {
-		if taskSemaphore.Acquire(ctx, 1) != nil {
+		if taskSemaphore.acquire(ctx) != nil {
 			break
 		}
 
 		wg.Add(1)
 		go func(sha1 string) {
-			defer taskSemaphore.Release(1)
+			defer taskSemaphore.release()
 			defer wg.Done()
 			if err := makeThumbnailFor(sha1); err != nil {
 				cancel(err)
@@ -1174,7 +1197,7 @@ func main() {
 		log.Fatalln("Unknown command: " + os.Args[1])
 	}
 
-	taskSemaphore = semaphore.NewWeighted(int64(runtime.NumCPU()))
+	taskSemaphore = newSemaphore(runtime.NumCPU())
 	err := cmd.handler(os.Args[2:])
 
 	// Note that the database object has a closing finalizer,