Repository: dragonflyoss/Dragonfly2
Pull request #1052 (2022-02-08 08:14:50)

fix: scheduler download tiny file error
Signed-off-by: Jim Ma <majinjing3@gmail.com>

<!--- Provide a general summary of your changes in the Title above -->

## Description

1. Fix wrong range when download tiny file.
2. Fix wrong peer id when download tiny file from dfget daemon.
3. Update unit tests.

## Related Issue

<!--- This project only accepts pull requests related to open issues -->
<!--- If suggesting a new feature or change, please discuss it in an issue first -->
<!--- If fixing a bug, there should be an issue describing it with steps to reproduce -->
<!--- Please link to the issue here: -->

## Motivation and Context

<!--- Why is this change required? What problem does it solve? -->

## Screenshots (if appropriate):

## Types of changes

<!--- What types of changes does your code introduce? Put an `x` in all the boxes that apply: -->

- [ ] Bug fix (non-breaking change which fixes an issue)
- [ ] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing functionality to change)
- [ ] Documentation Update (if none of the other choices apply)

## Checklist:

<!--- Go over all the following points, and put an `x` in all the boxes that apply. -->
<!--- If you're unsure about any of these, don't hesitate to ask. We're here to help! -->

- [ ] Code compiles correctly.
- [ ] My code follows the code style of this project.
- [ ] My change requires a change to the documentation.
- [ ] I have updated the documentation accordingly.
- [ ] I have read the **CONTRIBUTING** document.
- [ ] I have added tests to cover my changes.
- [ ] All new and existing tests passed.



Patch:
diff --git a/scheduler/resource/peer.go b/scheduler/resource/peer.go
index b5663386..a75440f8 100644
--- a/scheduler/resource/peer.go
+++ b/scheduler/resource/peer.go
@@ -385,25 +385,25 @@ func (p *Peer) DeleteStream() {
 	p.Stream = &atomic.Value{}
 }
 
-// Download tiny file from peer
+// DownloadTinyFile downloads tiny file from peer
 func (p *Peer) DownloadTinyFile() ([]byte, error) {
 	ctx, cancel := context.WithTimeout(context.Background(), downloadTinyFileContextTimeout)
 	defer cancel()
 
-	// Download url: http://${host}:${port}/download/${taskIndex}/${taskID}?peerId=scheduler
-	url := url.URL{
+	// Download url: http://${host}:${port}/download/${taskIndex}/${taskID}?peerId=${peerID}
+	targetURL := url.URL{
 		Scheme:   "http",
 		Host:     fmt.Sprintf("%s:%d", p.Host.IP, p.Host.DownloadPort),
 		Path:     fmt.Sprintf("download/%s/%s", p.Task.ID[:3], p.Task.ID),
-		RawQuery: "peerId=scheduler",
+		RawQuery: fmt.Sprintf("peerId=%s", p.ID),
 	}
 
-	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url.String(), nil)
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, targetURL.String(), nil)
 	if err != nil {
 		return []byte{}, err
 	}
-	req.Header.Set(headers.Range, fmt.Sprintf("bytes=%d-%d", 0, p.Task.ContentLength.Load()))
-	p.Log.Infof("download tiny file %s, header is : %#v", url.String(), req.Header)
+	req.Header.Set(headers.Range, fmt.Sprintf("bytes=%d-%d", 0, p.Task.ContentLength.Load()-1))
+	p.Log.Infof("download tiny file %s, header is : %#v", targetURL.String(), req.Header)
 
 	resp, err := http.DefaultClient.Do(req)
 	if err != nil {
@@ -415,7 +415,7 @@ func (p *Peer) DownloadTinyFile() ([]byte, error) {
 	// the request has succeeded and the body contains the requested ranges of data, as described in the Range header of the request.
 	// Refer to https://developer.mozilla.org/en-US/docs/Web/HTTP/Status/206
 	if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusBadRequest {
-		return []byte{}, fmt.Errorf("%v: %v", url.String(), resp.Status)
+		return []byte{}, fmt.Errorf("%v: %v", targetURL.String(), resp.Status)
 	}
 
 	return io.ReadAll(resp.Body)


Test patch:
diff --git a/scheduler/resource/peer_test.go b/scheduler/resource/peer_test.go
index cfc21cd7..b3052669 100644
--- a/scheduler/resource/peer_test.go
+++ b/scheduler/resource/peer_test.go
@@ -22,7 +22,6 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"net/url"
-	"path"
 	"strconv"
 	"testing"
 
@@ -30,6 +29,7 @@ import (
 	"github.com/golang/mock/gomock"
 	"github.com/stretchr/testify/assert"
 
+	"d7y.io/dragonfly/v2/client/clientutil"
 	"d7y.io/dragonfly/v2/pkg/idgen"
 	"d7y.io/dragonfly/v2/pkg/rpc/scheduler"
 	"d7y.io/dragonfly/v2/pkg/rpc/scheduler/mocks"
@@ -813,57 +813,79 @@ func TestPeer_DeleteStream(t *testing.T) {
 }
 
 func TestPeer_DownloadTinyFile(t *testing.T) {
-	s := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		if path.Base(r.URL.Path) == "foo" {
-			w.WriteHeader(http.StatusNotFound)
-			return
-		}
-
-		if r.Header.Get(headers.Range) == "bytes=0-2" {
-			w.WriteHeader(http.StatusNotAcceptable)
-			return
-		}
-
-		w.WriteHeader(http.StatusPartialContent)
-	}))
-	defer s.Close()
-
+	testData := []byte("./0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz" +
+		"./0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz")
+	newServer := func(t *testing.T, getPeer func() *Peer) *httptest.Server {
+		return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			peer := getPeer()
+			assert := assert.New(t)
+			assert.NotNil(peer)
+			assert.Equal(r.URL.Path, fmt.Sprintf("/download/%s/%s", peer.Task.ID[:3], peer.Task.ID))
+			assert.Equal(r.URL.RawQuery, fmt.Sprintf("peerId=%s", peer.ID))
+
+			rgs, err := clientutil.ParseRange(r.Header.Get(headers.Range), 128)
+			assert.Nil(err)
+			assert.Equal(1, len(rgs))
+			rg := rgs[0]
+
+			w.WriteHeader(http.StatusPartialContent)
+			n, err := w.Write(testData[rg.Start : rg.Start+rg.Length])
+			assert.Nil(err)
+			assert.Equal(int64(n), rg.Length)
+		}))
+	}
 	tests := []struct {
-		name   string
-		expect func(t *testing.T, peer *Peer)
+		name      string
+		newServer func(t *testing.T, getPeer func() *Peer) *httptest.Server
+		expect    func(t *testing.T, peer *Peer)
 	}{
 		{
-			name: "download tiny file",
+			name: "download tiny file - 32",
 			expect: func(t *testing.T, peer *Peer) {
 				assert := assert.New(t)
-				_, err := peer.DownloadTinyFile()
+				peer.Task.ContentLength.Store(32)
+				data, err := peer.DownloadTinyFile()
 				assert.NoError(err)
+				assert.Equal(testData[:32], data)
 			},
 		},
 		{
-			name: "download tiny file with range header",
+			name: "download tiny file - 128",
 			expect: func(t *testing.T, peer *Peer) {
 				assert := assert.New(t)
-				peer.Task.ContentLength.Store(2)
-				_, err := peer.DownloadTinyFile()
-				assert.EqualError(err, fmt.Sprintf("http://%s:%d/download/%s/%s?peerId=scheduler: 406 Not Acceptable",
-					peer.Host.IP, peer.Host.DownloadPort, peer.Task.ID[:3], peer.Task.ID))
+				peer.Task.ContentLength.Store(32)
+				data, err := peer.DownloadTinyFile()
+				assert.NoError(err)
+				assert.Equal(testData[:32], data)
 			},
 		},
 		{
 			name: "download tiny file failed because of http status code",
+			newServer: func(t *testing.T, getPeer func() *Peer) *httptest.Server {
+				return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+					w.WriteHeader(http.StatusNotFound)
+				}))
+			},
 			expect: func(t *testing.T, peer *Peer) {
 				assert := assert.New(t)
-				peer.Task.ID = "foo"
+				peer.Task.ID = "foobar"
 				_, err := peer.DownloadTinyFile()
-				assert.EqualError(err, fmt.Sprintf("http://%s:%d/download/%s/%s?peerId=scheduler: 404 Not Found",
-					peer.Host.IP, peer.Host.DownloadPort, peer.Task.ID[:3], peer.Task.ID))
+				assert.EqualError(err, fmt.Sprintf("http://%s:%d/download/%s/%s?peerId=%s: 404 Not Found",
+					peer.Host.IP, peer.Host.DownloadPort, peer.Task.ID[:3], peer.Task.ID, peer.ID))
 			},
 		},
 	}
 
 	for _, tc := range tests {
 		t.Run(tc.name, func(t *testing.T) {
+			var peer *Peer
+			if tc.newServer == nil {
+				tc.newServer = newServer
+			}
+			s := tc.newServer(t, func() *Peer {
+				return peer
+			})
+			defer s.Close()
 			url, err := url.Parse(s.URL)
 			if err != nil {
 				t.Fatal(err)
@@ -883,7 +905,7 @@ func TestPeer_DownloadTinyFile(t *testing.T) {
 			mockRawHost.DownPort = int32(port)
 			mockHost := NewHost(mockRawHost)
 			mockTask := NewTask(mockTaskID, mockTaskURL, mockTaskBackToSourceLimit, mockTaskURLMeta)
-			peer := NewPeer(mockPeerID, mockTask, mockHost)
+			peer = NewPeer(mockPeerID, mockTask, mockHost)
 			tc.expect(t, peer)
 		})
 	}