Merge pull request #5 from codex-storage/feat/add-context-cancellation

feat: context cancellation

Author: 2-towns

.gitignore

@@ -19,4 +19,7 @@ nimcache
 
 # Test files
 codex/testdata/hello.downloaded.txt
-codex/testdata/hello.downloaded.writer.txt
+codex/testdata/hello.downloaded.writer.txt
+
+# Bin
+codex-go

.vscode/settings.json

@@ -1,7 +1,9 @@
 {
     "go.toolsEnvVars": {
+        "CGO_ENABLED": "1",
         "CGO_CFLAGS": "-I${workspaceFolder}/vendor/nim-codex/library",
-        "CGO_LDFLAGS": "-L${workspaceFolder}/vendor/nim-codex/build -Wl,-rpath,${workspaceFolder}/vendor/nim-codex/build",
+        "CGO_LDFLAGS": "-L${workspaceFolder}/vendor/nim-codex/build -lcodex -Wl,-rpath,${workspaceFolder}/vendor/nim-codex/build",
         "LD_LIBRARY_PATH": "${workspaceFolder}/vendor/nim-codex/build:${env:LD_LIBRARY_PATH}"
-    }
+    },
+    "go.testTimeout": "2m"
 }

Makefile

@@ -29,7 +29,7 @@ build:
 
 test:
 	@echo "Running tests..."
-	CGO_ENABLED=1 CGO_CFLAGS="$(CGO_CFLAGS)" CGO_LDFLAGS="$(CGO_LDFLAGS)" go test ./...
+	CGO_ENABLED=1 CGO_CFLAGS="$(CGO_CFLAGS)" CGO_LDFLAGS="$(CGO_LDFLAGS)" GOTESTFLAGS="-timeout=2m" go test ./...
 
 clean:
 	@echo "Cleaning up..."

README.md

@@ -98,7 +98,7 @@ Now the module is ready for use in your project.
 
 The release process is defined [here](./RELEASE.md).
 
-## Usage
+## API
 
 ### Init
 
@@ -171,11 +171,10 @@ buf := bytes.NewBuffer([]byte("Hello World!"))
 onProgress := func(read, total int, percent float64, err error) {
    // Do something with the data
 }
-cid, err := codex.UploadReader(UploadOptions{filepath: "hello.txt", onProgress: onProgress}, buf)
+ctx := context.Background()
+cid, err := codex.UploadReader(ctx, UploadOptions{filepath: "hello.txt", onProgress: onProgress}, buf)
 ```
 
-Caveat: once started, the upload cannot be cancelled.
-
 #### file
 
 The `file` strategy allows you to upload a file on Codex using the path.
@@ -189,11 +188,10 @@ The `UploadFile` returns the cid of the content uploaded.
 onProgress := func(read, total int, percent float64, err error) {
    // Do something with the data
 }
-cid, err := codex.UploadFile(UploadOptions{filepath: "./testdata/hello.txt", onProgress: onProgress})
+ctx := context.Background()
+cid, err := codex.UploadFile(ctx, UploadOptions{filepath: "./testdata/hello.txt", onProgress: onProgress})
 ```
 
-Caveat: once started, the upload cannot be cancelled.
-
 #### chunks
 
 The `chunks` strategy allows you to manage the upload by yourself. It requires more code
@@ -246,11 +244,10 @@ opt := DownloadStreamOptions{
       // Handle progress
    },
 }
-err := codex.DownloadStream(cid, opt)
+ctx := context.Background()
+err := codex.DownloadStream(ctx, cid, opt)
 ```
 
-Caveat: once started, the download cannot be cancelled.
-
 #### chunks
 
 The `chunks` strategy allows to manage the download by yourself. It requires more code
@@ -310,4 +307,10 @@ peerId := "..."
 record, err := node.CodexPeerDebug(peerId)
 ```
 
-`CodexPeerDebug` is only available if you built with `-d:codex_enable_api_debug_peers=true` flag.
+`CodexPeerDebug` is only available if you built with `-d:codex_enable_api_debug_peers=true` flag.
+
+### Context and cancellation
+
+Go contexts are exposed only on the long-running operations as `UploadReader`, `UploadFile`, and `DownloadFile`. If the
+context is cancelled, those methods cancel the active upload or download. Short lived API calls don’t take a context
+because they usually finish before a cancellation signal could matter.

codex/codex.go

@@ -92,7 +92,7 @@ const (
 
 type Config struct {
 	// Default: INFO
-	LogLevel LogLevel `json:"log-level,omitempty"`
+	LogLevel string `json:"log-level,omitempty"`
 
 	// Specifies what kind of logs should be written to stdout
 	// Default: auto
@@ -280,8 +280,12 @@ func (node CodexNode) Destroy() error {
 		return bridge.callError("cGoCodexDestroy")
 	}
 
-	_, err = bridge.wait()
-	return err
+	// We don't wait for the bridge here.
+	// The destroy function does not call the worker thread,
+	// it destroys the context directly and return the return
+	// value synchronously.
+
+	return nil
 }
 
 // Version returns the version of the Codex node.

codex/codex_test.go

@@ -3,7 +3,11 @@ package codex
 import "testing"
 
 func TestCodexVersion(t *testing.T) {
-	node := newCodexNode(t, withNoStart())
+	config := defaultConfigHelper(t)
+	node, err := New(config)
+	if err != nil {
+		t.Fatalf("Failed to create Codex node: %v", err)
+	}
 
 	version, err := node.Version()
 	if err != nil {
@@ -17,7 +21,11 @@ func TestCodexVersion(t *testing.T) {
 }
 
 func TestCodexRevision(t *testing.T) {
-	node := newCodexNode(t, withNoStart())
+	config := defaultConfigHelper(t)
+	node, err := New(config)
+	if err != nil {
+		t.Fatalf("Failed to create Codex node: %v", err)
+	}
 
 	revision, err := node.Revision()
 	if err != nil {

codex/debug_test.go

@@ -32,45 +32,31 @@ func TestUpdateLogLevel(t *testing.T) {
 	}
 	defer os.Remove(tmpFile.Name())
 
-	node, err := New(Config{
-		LogFile:        tmpFile.Name(),
-		MetricsEnabled: false,
+	node := newCodexNode(t, Config{
+		LogLevel:  "INFO",
+		LogFile:   tmpFile.Name(),
+		LogFormat: LogFormatNoColors,
 	})
-	if err != nil {
-		t.Fatalf("Failed to create Codex node: %v", err)
-	}
-
-	t.Cleanup(func() {
-		if err := node.Stop(); err != nil {
-			t.Logf("cleanup codex: %v", err)
-		}
-
-		if err := node.Destroy(); err != nil {
-			t.Logf("cleanup codex: %v", err)
-		}
-	})
-
-	if err := node.Start(); err != nil {
-		t.Fatalf("Failed to start Codex node: %v", err)
-	}
 
 	content, err := os.ReadFile(tmpFile.Name())
+
 	if err != nil {
 		t.Fatalf("Failed to read log file: %v", err)
 	}
-	if !strings.Contains(string(content), "Started codex node") {
-		t.Errorf("Log file does not contain 'Started codex node' %s", string(content))
-	}
-
-	if err := node.Stop(); err != nil {
-		t.Fatalf("Failed to stop Codex node: %v", err)
+	if !strings.Contains(string(content), "INF") {
+		t.Errorf("Log file does not contain INFO statement %s", string(content))
 	}
 
 	err = node.UpdateLogLevel("ERROR")
 	if err != nil {
 		t.Fatalf("UpdateLogLevel call failed: %v", err)
 	}
 
+	if err := node.Stop(); err != nil {
+		t.Fatalf("Failed to stop Codex node: %v", err)
+	}
+
+	// Clear the file
 	if err := os.WriteFile(tmpFile.Name(), []byte{}, 0644); err != nil {
 		t.Fatalf("Failed to clear log file: %v", err)
 	}
@@ -85,58 +71,18 @@ func TestUpdateLogLevel(t *testing.T) {
 		t.Fatalf("Failed to read log file: %v", err)
 	}
 
-	if strings.Contains(string(content), "Starting discovery node") {
-		t.Errorf("Log file contains 'Starting discovery node'")
+	if strings.Contains(string(content), "INF") {
+		t.Errorf("Log file contains INFO statement after log level update: %s", string(content))
 	}
 }
 
 func TestCodexPeerDebug(t *testing.T) {
 	var bootstrap, node1, node2 *CodexNode
 	var err error
 
-	t.Cleanup(func() {
-		if bootstrap != nil {
-			if err := bootstrap.Stop(); err != nil {
-				t.Logf("cleanup bootstrap: %v", err)
-			}
-
-			if err := bootstrap.Destroy(); err != nil {
-				t.Logf("cleanup bootstrap: %v", err)
-			}
-		}
-		if node1 != nil {
-			if err := node1.Stop(); err != nil {
-				t.Logf("cleanup node1: %v", err)
-			}
-
-			if err := node1.Destroy(); err != nil {
-				t.Logf("cleanup node1: %v", err)
-			}
-		}
-		if node2 != nil {
-			if err := node2.Stop(); err != nil {
-				t.Logf("cleanup node2: %v", err)
-			}
-
-			if err := node2.Destroy(); err != nil {
-				t.Logf("cleanup node2: %v", err)
-			}
-		}
-	})
-
-	bootstrap, err = New(Config{
-		DataDir:        t.TempDir(),
-		LogFormat:      LogFormatNoColors,
-		MetricsEnabled: false,
-		DiscoveryPort:  8092,
+	bootstrap = newCodexNode(t, Config{
+		DiscoveryPort: 8092,
 	})
-	if err != nil {
-		t.Fatalf("Failed to create bootstrap: %v", err)
-	}
-
-	if err := bootstrap.Start(); err != nil {
-		t.Fatalf("Failed to start bootstrap: %v", err)
-	}
 
 	spr, err := bootstrap.Spr()
 	if err != nil {
@@ -145,35 +91,15 @@ func TestCodexPeerDebug(t *testing.T) {
 
 	bootstrapNodes := []string{spr}
 
-	node1, err = New(Config{
-		DataDir:        t.TempDir(),
-		LogFormat:      LogFormatNoColors,
-		MetricsEnabled: false,
+	node1 = newCodexNode(t, Config{
 		DiscoveryPort:  8090,
 		BootstrapNodes: bootstrapNodes,
 	})
-	if err != nil {
-		t.Fatalf("Failed to create codex: %v", err)
-	}
 
-	if err := node1.Start(); err != nil {
-		t.Fatalf("Failed to start codex: %v", err)
-	}
-
-	node2, err = New(Config{
-		DataDir:        t.TempDir(),
-		LogFormat:      LogFormatNoColors,
-		MetricsEnabled: false,
+	node2 = newCodexNode(t, Config{
 		DiscoveryPort:  8091,
 		BootstrapNodes: bootstrapNodes,
 	})
-	if err != nil {
-		t.Fatalf("Failed to create codex2: %v", err)
-	}
-
-	if err := node2.Start(); err != nil {
-		t.Fatalf("Failed to start codex2: %v", err)
-	}
 
 	peerId, err := node2.PeerId()
 	if err != nil {
@@ -186,9 +112,14 @@ func TestCodexPeerDebug(t *testing.T) {
 		if err == nil {
 			break
 		}
+
 		time.Sleep(1 * time.Second)
 	}
 
+	if err != nil {
+		t.Fatalf("CodexPeerDebug call failed: %v", err)
+	}
+
 	if record.PeerId == "" {
 		t.Fatalf("CodexPeerDebug call failed: %v", err)
 	}

codex/download.go

@@ -26,7 +26,9 @@ package codex
 */
 import "C"
 import (
+	"context"
 	"encoding/json"
+	"fmt"
 	"io"
 	"unsafe"
 )
@@ -145,7 +147,7 @@ func (node CodexNode) DownloadManifest(cid string) (Manifest, error) {
 // If options.writer is set, the data will be written into that writer.
 // The options filepath and writer are not mutually exclusive, i.e you can write
 // in different places in a same call.
-func (node CodexNode) DownloadStream(cid string, options DownloadStreamOptions) error {
+func (node CodexNode) DownloadStream(ctx context.Context, cid string, options DownloadStreamOptions) error {
 	bridge := newBridgeCtx()
 	defer bridge.free()
 
@@ -189,6 +191,16 @@ func (node CodexNode) DownloadStream(cid string, options DownloadStreamOptions)
 	var cCid = C.CString(cid)
 	defer C.free(unsafe.Pointer(cCid))
 
+	err := node.DownloadInit(cid, DownloadInitOptions{
+		ChunkSize: options.ChunkSize,
+		Local:     options.Local,
+	})
+	if err != nil {
+		return err
+	}
+
+	defer node.DownloadCancel(cid)
+
 	var cFilepath = C.CString(options.Filepath)
 	defer C.free(unsafe.Pointer(cFilepath))
 
@@ -198,8 +210,39 @@ func (node CodexNode) DownloadStream(cid string, options DownloadStreamOptions)
 		return bridge.callError("cGoCodexDownloadLocal")
 	}
 
-	_, err := bridge.wait()
-	return err
+	// Create a done channel to signal the goroutine to stop
+	// when the download is complete and avoid goroutine leaks.
+	done := make(chan struct{})
+	defer close(done)
+
+	channelError := make(chan error, 1)
+	go func() {
+		select {
+		case <-ctx.Done():
+			channelError <- node.DownloadCancel(cid)
+		case <-done:
+			// Nothing to do, download finished
+		}
+	}()
+
+	_, err = bridge.wait()
+
+	// Extract the potential cancellation error
+	var cancelError error
+	select {
+	case cancelError = <-channelError:
+	default:
+	}
+
+	if err != nil {
+		if cancelError != nil {
+			return fmt.Errorf("download canceled: %v, but failed to cancel download session: %v", ctx.Err(), cancelError)
+		}
+
+		return err
+	}
+
+	return cancelError
 }
 
 // DownloadInit initializes the download process for a specific CID.