nip13: parallel DoWork(), deprecate Generate()

nbd-wtf · Aug 21, 2024 · 83bd196 · 83bd196
1 parent 2da5835
commit 83bd196
Show file tree

Hide file tree

Showing 2 changed files with 90 additions and 41 deletions.
diff --git a/nip13/nip13.go b/nip13/nip13.go
@@ -1,9 +1,11 @@
 package nip13
 
 import (
+	"context"
 	"encoding/hex"
 	"errors"
 	"math/bits"
+	"runtime"
 	"strconv"
 	"time"
 
@@ -63,12 +65,7 @@ func Check(id string, minDifficulty int) error {
 	return nil
 }
 
-// Generate performs proof of work on the specified event until either the target
-// difficulty is reached or the function runs for longer than the timeout.
-// The latter case results in ErrGenerateTimeout.
-//
-// Upon success, the returned event always contains a "nonce" tag with the target difficulty
-// commitment, and an updated event.CreatedAt.
+// Deprecated: use DoWork()
 func Generate(event *nostr.Event, targetDifficulty int, timeout time.Duration) (*nostr.Event, error) {
 	if event.PubKey == "" {
 		return nil, ErrMissingPubKey
@@ -91,3 +88,53 @@ func Generate(event *nostr.Event, targetDifficulty int, timeout time.Duration) (
 		}
 	}
 }
+
+// DoWork() performs work in multiple threads (given by runtime.NumCPU()) and returns the first
+// nonce (as a nostr.Tag) that yields the required work.
+// Returns an error if the context expires before that.
+func DoWork(ctx context.Context, event nostr.Event, targetDifficulty int) (nostr.Tag, error) {
+	if event.PubKey == "" {
+		return nil, ErrMissingPubKey
+	}
+
+	ctx, cancel := context.WithCancel(ctx)
+	nthreads := runtime.NumCPU()
+	var nonceTag nostr.Tag
+
+	for i := 0; i < nthreads; i++ {
+		go func(event nostr.Event, nonce uint64) {
+			tag := nostr.Tag{"nonce", "", strconv.Itoa(targetDifficulty)}
+			event.Tags = append(event.Tags, tag)
+			for {
+				// try 10000 times (~30ms)
+				for n := 0; n < 10000; n++ {
+					tag[1] = strconv.FormatUint(nonce, 10)
+
+					if Difficulty(event.GetID()) >= targetDifficulty {
+						nonceTag = tag
+						cancel()
+						return
+					}
+
+					nonce += uint64(nthreads)
+				}
+
+				// then check if the context was canceled
+				select {
+				case <-ctx.Done():
+					return
+				default:
+					// otherwise keep trying
+				}
+			}
+		}(event, uint64(i))
+	}
+
+	<-ctx.Done()
+
+	if nonceTag != nil {
+		return nonceTag, nil
+	}
+
+	return nil, ErrGenerateTimeout
+}
diff --git a/nip13/nip13_test.go b/nip13/nip13_test.go
@@ -1,6 +1,7 @@
 package nip13
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"strconv"
@@ -52,51 +53,48 @@ func TestCommittedDifficulty(t *testing.T) {
 	}
 }
 
-func TestGenerateShort(t *testing.T) {
-	event := &nostr.Event{
+func TestDoWorkShort(t *testing.T) {
+	event := nostr.Event{
 		Kind:    nostr.KindTextNote,
 		Content: "It's just me mining my own business",
 		PubKey:  "a48380f4cfcc1ad5378294fcac36439770f9c878dd880ffa94bb74ea54a6f243",
 	}
-	pow, err := Generate(event, 0, 3*time.Second)
+	pow, err := DoWork(context.Background(), event, 0)
 	if err != nil {
 		t.Fatal(err)
 	}
 	testNonceTag(t, pow, 0)
 }
 
-func TestGenerateLong(t *testing.T) {
+func TestDoWorkLong(t *testing.T) {
 	if testing.Short() {
 		t.Skip("too consuming for short mode")
 	}
 	for _, difficulty := range []int{8, 16} {
 		difficulty := difficulty
 		t.Run(fmt.Sprintf("%dbits", difficulty), func(t *testing.T) {
 			t.Parallel()
-			event := &nostr.Event{
+			event := nostr.Event{
 				Kind:    nostr.KindTextNote,
 				Content: "It's just me mining my own business",
 				PubKey:  "a48380f4cfcc1ad5378294fcac36439770f9c878dd880ffa94bb74ea54a6f243",
 			}
-			pow, err := Generate(event, difficulty, time.Minute)
+			ctx, _ := context.WithTimeout(context.Background(), time.Minute)
+			pow, err := DoWork(ctx, event, difficulty)
 			if err != nil {
 				t.Fatal(err)
 			}
-			if err := Check(pow.GetID(), difficulty); err != nil {
+			event.Tags = append(event.Tags, pow)
+			if err := Check(event.GetID(), difficulty); err != nil {
 				t.Error(err)
 			}
 			testNonceTag(t, pow, difficulty)
 		})
 	}
 }
 
-func testNonceTag(t *testing.T, event *nostr.Event, commitment int) {
+func testNonceTag(t *testing.T, tag nostr.Tag, commitment int) {
 	t.Helper()
-	tagptr := event.Tags.GetFirst([]string{"nonce"})
-	if tagptr == nil {
-		t.Fatal("no nonce tag")
-	}
-	tag := *tagptr
 	if tag[0] != "nonce" {
 		t.Errorf("tag[0] = %q; want 'nonce'", tag[0])
 	}
@@ -108,47 +106,50 @@ func testNonceTag(t *testing.T, event *nostr.Event, commitment int) {
 	}
 }
 
-func TestGenerateTimeout(t *testing.T) {
-	event := &nostr.Event{
+func TestDoWorkTimeout(t *testing.T) {
+	event := nostr.Event{
 		Kind:    nostr.KindTextNote,
 		Content: "It's just me mining my own business",
 		PubKey:  "a48380f4cfcc1ad5378294fcac36439770f9c878dd880ffa94bb74ea54a6f243",
 	}
 	done := make(chan error)
 	go func() {
-		_, err := Generate(event, 256, time.Millisecond)
+		ctx, _ := context.WithTimeout(context.Background(), time.Millisecond)
+		_, err := DoWork(ctx, event, 256)
 		done <- err
 	}()
 	select {
 	case <-time.After(time.Second):
-		t.Error("Generate took too long to timeout")
+		t.Error("DoWork took too long to timeout")
 	case err := <-done:
 		if !errors.Is(err, ErrGenerateTimeout) {
-			t.Errorf("Generate returned %v; want ErrGenerateTimeout", err)
+			t.Errorf("DoWork returned %v; want ErrDoWorkTimeout", err)
 		}
 	}
 }
 
-func BenchmarkCheck(b *testing.B) {
-	for i := 0; i < b.N; i++ {
-		Check("000000000e9d97a1ab09fc381030b346cdd7a142ad57e6df0b46dc9bef6c7e2d", 36)
+func BenchmarkGenerate(b *testing.B) {
+	if testing.Short() {
+		b.Skip("too consuming for short mode")
 	}
-}
-
-func BenchmarkGenerateOneIteration(b *testing.B) {
-	for i := 0; i < b.N; i++ {
-		event := &nostr.Event{
-			Kind:    nostr.KindTextNote,
-			Content: "It's just me mining my own business",
-			PubKey:  "a48380f4cfcc1ad5378294fcac36439770f9c878dd880ffa94bb74ea54a6f243",
-		}
-		if _, err := Generate(event, 0, time.Minute); err != nil {
-			b.Fatal(err)
-		}
+	for _, difficulty := range []int{8, 16, 24} {
+		difficulty := difficulty
+		b.Run(fmt.Sprintf("%dbits", difficulty), func(b *testing.B) {
+			for i := 0; i < b.N; i++ {
+				event := &nostr.Event{
+					Kind:    nostr.KindTextNote,
+					Content: "It's just me mining my own business",
+					PubKey:  "a48380f4cfcc1ad5378294fcac36439770f9c878dd880ffa94bb74ea54a6f243",
+				}
+				if _, err := Generate(event, difficulty, time.Minute); err != nil {
+					b.Fatal(err)
+				}
+			}
+		})
 	}
 }
 
-func BenchmarkGenerate(b *testing.B) {
+func BenchmarkDoWork(b *testing.B) {
 	if testing.Short() {
 		b.Skip("too consuming for short mode")
 	}
@@ -161,7 +162,8 @@ func BenchmarkGenerate(b *testing.B) {
 					Content: "It's just me mining my own business",
 					PubKey:  "a48380f4cfcc1ad5378294fcac36439770f9c878dd880ffa94bb74ea54a6f243",
 				}
-				if _, err := Generate(event, difficulty, time.Minute); err != nil {
+				ctx, _ := context.WithTimeout(context.Background(), time.Second*60)
+				if _, err := DoWork(ctx, *event, difficulty); err != nil {
 					b.Fatal(err)
 				}
 			}