Repository: dragonflyoss/Dragonfly2
Pull request #1024 (2022-01-21 08:04:05)

feat: scheduler download tiny file with range header
Signed-off-by: Gaius <gaius.qi@gmail.com>

<!--- Provide a general summary of your changes in the Title above -->

## Description
- Scheduler download tiny file with range header and handle http error code.
<!--- Describe your changes in detail -->

## 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)
- [x] 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! -->

- [x] Code compiles correctly.
- [x] My code follows the code style of this project.
- [ ] My change requires a change to the documentation.
- [ ] I have updated the documentation accordingly.
- [x] I have read the **CONTRIBUTING** document.
- [x] I have added tests to cover my changes.
- [x] All new and existing tests passed.


codecov-commenter: # [Codecov](https://codecov.io/gh/dragonflyoss/Dragonfly2/pull/1024?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=dragonflyoss) Report
> Merging [#1024](https://codecov.io/gh/dragonflyoss/Dragonfly2/pull/1024?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=dragonflyoss) (9ee832d) into [main](https://codecov.io/gh/dragonflyoss/Dragonfly2/commit/037b77df3c5954acfb4f6703030b5ff28d9d103a?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=dragonflyoss) (037b77d) will **increase** coverage by `0.00%`.
> The diff coverage is `100.00%`.

[![Impacted file tree graph](https://codecov.io/gh/dragonflyoss/Dragonfly2/pull/1024/graphs/tree.svg?width=650&height=150&src=pr&token=PUICJYSAT7&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=dragonflyoss)](https://codecov.io/gh/dragonflyoss/Dragonfly2/pull/1024?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=dragonflyoss)

```diff
@@           Coverage Diff           @@
##             main    #1024   +/-   ##
=======================================
  Coverage   45.96%   45.96%           
=======================================
  Files         128      128           
  Lines        9134     9142    +8     
=======================================
+ Hits         4198     4202    +4     
- Misses       4463     4467    +4     
  Partials      473      473           
```

| Flag | Coverage Δ | |
|---|---|---|
| Object-compatibility-e2etests | `?` | |
| e2etests | `?` | |
| unittests | `45.96% <100.00%> (+<0.01%)` | :arrow_up: |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=dragonflyoss#carryforward-flags-in-the-pull-request-comment) to find out more.

| [Impacted Files](https://codecov.io/gh/dragonflyoss/Dragonfly2/pull/1024?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=dragonflyoss) | Coverage Δ | |
|---|---|---|
| [scheduler/resource/peer.go](https://codecov.io/gh/dragonflyoss/Dragonfly2/pull/1024/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=dragonflyoss#diff-c2NoZWR1bGVyL3Jlc291cmNlL3BlZXIuZ28=) | `74.84% <100.00%> (+0.48%)` | :arrow_up: |
| [client/daemon/transport/transport.go](https://codecov.io/gh/dragonflyoss/Dragonfly2/pull/1024/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=dragonflyoss#diff-Y2xpZW50L2RhZW1vbi90cmFuc3BvcnQvdHJhbnNwb3J0Lmdv) | `57.98% <0.00%> (-1.50%)` | :arrow_down: |
| [client/daemon/peer/peertask\_conductor.go](https://codecov.io/gh/dragonflyoss/Dragonfly2/pull/1024/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=dragonflyoss#diff-Y2xpZW50L2RhZW1vbi9wZWVyL3BlZXJ0YXNrX2NvbmR1Y3Rvci5nbw==) | `56.91% <0.00%> (-0.42%)` | :arrow_down: |
| [client/daemon/proxy/proxy.go](https://codecov.io/gh/dragonflyoss/Dragonfly2/pull/1024/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=dragonflyoss#diff-Y2xpZW50L2RhZW1vbi9wcm94eS9wcm94eS5nbw==) | `9.74% <0.00%> (-0.04%)` | :arrow_down: |
| [client/daemon/peer/peertask\_manager.go](https://codecov.io/gh/dragonflyoss/Dragonfly2/pull/1024/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=dragonflyoss#diff-Y2xpZW50L2RhZW1vbi9wZWVyL3BlZXJ0YXNrX21hbmFnZXIuZ28=) | `44.30% <0.00%> (+0.14%)` | :arrow_up: |
| [cdn/supervisor/progress/progress.go](https://codecov.io/gh/dragonflyoss/Dragonfly2/pull/1024/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=dragonflyoss#diff-Y2RuL3N1cGVydmlzb3IvcHJvZ3Jlc3MvcHJvZ3Jlc3MuZ28=) | `97.01% <0.00%> (+5.97%)` | :arrow_up: |


Patch:
diff --git a/scheduler/resource/peer.go b/scheduler/resource/peer.go
index 8d4f5406..eb214b09 100644
--- a/scheduler/resource/peer.go
+++ b/scheduler/resource/peer.go
@@ -26,6 +26,7 @@ import (
 	"time"
 
 	"github.com/bits-and-blooms/bitset"
+	"github.com/go-http-utils/headers"
 	"github.com/looplab/fsm"
 	"go.uber.org/atomic"
 
@@ -402,6 +403,7 @@ func (p *Peer) DownloadTinyFile(ctx context.Context) ([]byte, error) {
 	if err != nil {
 		return []byte{}, err
 	}
+	req.Header.Set(headers.Range, fmt.Sprintf("bytes=%d-%d", 0, p.Task.ContentLength.Load()))
 
 	resp, err := http.DefaultClient.Do(req)
 	if err != nil {
@@ -409,5 +411,9 @@ func (p *Peer) DownloadTinyFile(ctx context.Context) ([]byte, error) {
 	}
 	defer resp.Body.Close()
 
+	if resp.StatusCode != http.StatusOK {
+		return []byte{}, fmt.Errorf("%v: %v", url.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 73ea59d8..77466dc9 100644
--- a/scheduler/resource/peer_test.go
+++ b/scheduler/resource/peer_test.go
@@ -18,14 +18,17 @@ package resource
 
 import (
 	"context"
+	"fmt"
 	"net"
 	"net/http"
 	"net/http/httptest"
 	"net/url"
+	"path"
 	"strconv"
 	"testing"
 	"time"
 
+	"github.com/go-http-utils/headers"
 	"github.com/golang/mock/gomock"
 	"github.com/stretchr/testify/assert"
 
@@ -867,6 +870,16 @@ 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.StatusOK)
 	}))
 	defer s.Close()
@@ -886,6 +899,19 @@ func TestPeer_DownloadTinyFile(t *testing.T) {
 				assert.NoError(err)
 			},
 		},
+		{
+			name: "download tiny file with range header",
+			expect: func(t *testing.T, peer *Peer) {
+				assert := assert.New(t)
+				peer.Task.ContentLength.Store(2)
+				ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
+				defer cancel()
+
+				_, err := peer.DownloadTinyFile(ctx)
+				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))
+			},
+		},
 		{
 			name: "download tiny file failed because of port error",
 			expect: func(t *testing.T, peer *Peer) {
@@ -910,6 +936,19 @@ func TestPeer_DownloadTinyFile(t *testing.T) {
 				assert.Error(err)
 			},
 		},
+		{
+			name: "download tiny file failed because of http status code",
+			expect: func(t *testing.T, peer *Peer) {
+				assert := assert.New(t)
+				peer.Task.ID = "foo"
+				ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
+				defer cancel()
+
+				_, err := peer.DownloadTinyFile(ctx)
+				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))
+			},
+		},
 	}
 
 	for _, tc := range tests {