Repository: dragonflyoss/Dragonfly2
Pull request #1557 (2022-08-12 04:12:03)

feat: storage calculate the count of records
Signed-off-by: Gaius <gaius.qi@gmail.com>

<!--- Provide a general summary of your changes in the Title above -->

## Description
- Storage calculate the count of records.
<!--- 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! -->

- [ ] My change requires a change to the documentation.
- [ ] I have updated the documentation accordingly.
- [x] I have read the **CONTRIBUTING** document.
- [ ] I have added tests to cover my changes.


codecov-commenter: # [Codecov](https://codecov.io/gh/dragonflyoss/Dragonfly2/pull/1557?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=dragonflyoss) Report
> Merging [#1557](https://codecov.io/gh/dragonflyoss/Dragonfly2/pull/1557?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=dragonflyoss) (c3cbae2) into [main](https://codecov.io/gh/dragonflyoss/Dragonfly2/commit/fbb6f1d081b4860760640cf345d0286c8b9c51f9?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=dragonflyoss) (fbb6f1d) will **decrease** coverage by `0.12%`.
> The diff coverage is `100.00%`.

[![Impacted file tree graph](https://codecov.io/gh/dragonflyoss/Dragonfly2/pull/1557/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/1557?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=dragonflyoss)

```diff
@@            Coverage Diff             @@
##             main    #1557      +/-   ##
==========================================
- Coverage   43.29%   43.16%   -0.13%     
==========================================
  Files         118      118              
  Lines       10919    10925       +6     
==========================================
- Hits         4727     4716      -11     
- Misses       5708     5720      +12     
- Partials      484      489       +5     
```

| Flag | Coverage Δ | |
|---|---|---|
| Object-compatibility-e2etests | `?` | |
| e2etests | `?` | |
| unittests | `43.16% <100.00%> (-0.13%)` | :arrow_down: |

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/1557?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=dragonflyoss) | Coverage Δ | |
|---|---|---|
| [scheduler/storage/storage.go](https://codecov.io/gh/dragonflyoss/Dragonfly2/pull/1557/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=dragonflyoss#diff-c2NoZWR1bGVyL3N0b3JhZ2Uvc3RvcmFnZS5nbw==) | `82.30% <100.00%> (-0.88%)` | :arrow_down: |
| [...ent/daemon/peer/peertask\_piecetask\_synchronizer.go](https://codecov.io/gh/dragonflyoss/Dragonfly2/pull/1557/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=dragonflyoss#diff-Y2xpZW50L2RhZW1vbi9wZWVyL3BlZXJ0YXNrX3BpZWNldGFza19zeW5jaHJvbml6ZXIuZ28=) | `47.61% <0.00%> (-1.30%)` | :arrow_down: |
| [client/daemon/peer/peertask\_conductor.go](https://codecov.io/gh/dragonflyoss/Dragonfly2/pull/1557/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=dragonflyoss#diff-Y2xpZW50L2RhZW1vbi9wZWVyL3BlZXJ0YXNrX2NvbmR1Y3Rvci5nbw==) | `56.13% <0.00%> (-0.87%)` | :arrow_down: |
| [client/daemon/peer/piece\_manager.go](https://codecov.io/gh/dragonflyoss/Dragonfly2/pull/1557/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=dragonflyoss#diff-Y2xpZW50L2RhZW1vbi9wZWVyL3BpZWNlX21hbmFnZXIuZ28=) | `53.01% <0.00%> (-0.54%)` | :arrow_down: |


Patch:
diff --git a/manager/manager.go b/manager/manager.go
index 718e5c03..c7e0ac32 100644
--- a/manager/manager.go
+++ b/manager/manager.go
@@ -219,15 +219,17 @@ func (s *Server) Stop() {
 	// Stop REST server
 	if err := s.restServer.Shutdown(context.Background()); err != nil {
 		logger.Errorf("rest server failed to stop: %+v", err)
+	} else {
+		logger.Info("rest server closed under request")
 	}
-	logger.Info("rest server closed under request")
 
 	// Stop metrics server
 	if s.metricsServer != nil {
 		if err := s.metricsServer.Shutdown(context.Background()); err != nil {
 			logger.Errorf("metrics server failed to stop: %+v", err)
+		} else {
+			logger.Info("metrics server closed under request")
 		}
-		logger.Info("metrics server closed under request")
 	}
 
 	// Stop GRPC server
diff --git a/scheduler/scheduler.go b/scheduler/scheduler.go
index 5e98372c..fef51ffe 100644
--- a/scheduler/scheduler.go
+++ b/scheduler/scheduler.go
@@ -67,6 +67,9 @@ type Server struct {
 	// Async job.
 	job job.Job
 
+	// Storage service.
+	storage storage.Storage
+
 	// GC server.
 	gc gc.GC
 }
@@ -135,9 +138,10 @@ func New(ctx context.Context, cfg *config.Config, d dfpath.Dfpath) (*Server, err
 	if err != nil {
 		return nil, err
 	}
+	s.storage = storage
 
 	// Initialize scheduler service.
-	service := service.New(cfg, resource, scheduler, dynconfig, storage)
+	service := service.New(cfg, resource, scheduler, dynconfig, s.storage)
 
 	// Initialize grpc service.
 	var schedulerServerOptions []grpc.ServerOption
@@ -234,15 +238,24 @@ func (s *Server) Stop() {
 	// Stop dynconfig server.
 	if err := s.dynconfig.Stop(); err != nil {
 		logger.Errorf("dynconfig client closed failed %s", err.Error())
+	} else {
+		logger.Info("dynconfig client closed")
 	}
-	logger.Info("dynconfig client closed")
 
 	// Stop manager client.
 	if s.managerClient != nil {
 		if err := s.managerClient.Close(); err != nil {
 			logger.Errorf("manager client failed to stop: %s", err.Error())
+		} else {
+			logger.Info("manager client closed")
 		}
-		logger.Info("manager client closed")
+	}
+
+	// Clean storage.
+	if err := s.storage.Clear(); err != nil {
+		logger.Errorf("clean storage failed %s", err.Error())
+	} else {
+		logger.Info("clean storage completed")
 	}
 
 	// Stop GC.
@@ -253,8 +266,9 @@ func (s *Server) Stop() {
 	if s.metricsServer != nil {
 		if err := s.metricsServer.Shutdown(context.Background()); err != nil {
 			logger.Errorf("metrics server failed to stop: %s", err.Error())
+		} else {
+			logger.Info("metrics server closed under request")
 		}
-		logger.Info("metrics server closed under request")
 	}
 
 	// Stop GRPC server.
diff --git a/scheduler/storage/mocks/storage_mock.go b/scheduler/storage/mocks/storage_mock.go
index 412bf0d0..129d6884 100644
--- a/scheduler/storage/mocks/storage_mock.go
+++ b/scheduler/storage/mocks/storage_mock.go
@@ -49,6 +49,20 @@ func (mr *MockStorageMockRecorder) Clear() *gomock.Call {
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Clear", reflect.TypeOf((*MockStorage)(nil).Clear))
 }
 
+// Count mocks base method.
+func (m *MockStorage) Count() int64 {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Count")
+	ret0, _ := ret[0].(int64)
+	return ret0
+}
+
+// Count indicates an expected call of Count.
+func (mr *MockStorageMockRecorder) Count() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Count", reflect.TypeOf((*MockStorage)(nil).Count))
+}
+
 // Create mocks base method.
 func (m *MockStorage) Create(arg0 storage.Record) error {
 	m.ctrl.T.Helper()
diff --git a/scheduler/storage/storage.go b/scheduler/storage/storage.go
index 1f64466a..1beb9e30 100644
--- a/scheduler/storage/storage.go
+++ b/scheduler/storage/storage.go
@@ -183,6 +183,9 @@ type Storage interface {
 	// List returns all of records in csv file.
 	List() ([]Record, error)
 
+	// Count returns the count of records.
+	Count() int64
+
 	// Open opens storage for read, it returns io.ReadCloser of storage files.
 	Open() (io.ReadCloser, error)
 
@@ -198,6 +201,7 @@ type storage struct {
 	maxBackups int
 	buffer     []Record
 	bufferSize int
+	count      int64
 	mu         *sync.RWMutex
 }
 
@@ -257,12 +261,26 @@ func (s *storage) Create(record Record) error {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
+	// Write without buffer.
+	if s.bufferSize == 0 {
+		if err := s.create(s.buffer...); err != nil {
+			return err
+		}
+
+		// Update record count.
+		s.count++
+		return nil
+	}
+
 	// Write records to file.
 	if len(s.buffer) >= s.bufferSize {
 		if err := s.create(s.buffer...); err != nil {
 			return err
 		}
 
+		// Update record count.
+		s.count += int64(s.bufferSize)
+
 		// Keep allocated memory.
 		s.buffer = s.buffer[:0]
 	}
@@ -310,6 +328,11 @@ func (s *storage) List() ([]Record, error) {
 	return records, nil
 }
 
+// Count returns the count of records.
+func (s *storage) Count() int64 {
+	return s.count
+}
+
 // Open opens storage for read, it returns io.ReadCloser of storage files.
 func (s *storage) Open() (io.ReadCloser, error) {
 	s.mu.RLock()


Test patch:
diff --git a/scheduler/storage/storage_test.go b/scheduler/storage/storage_test.go
index 56e1f1a0..24e23897 100644
--- a/scheduler/storage/storage_test.go
+++ b/scheduler/storage/storage_test.go
@@ -51,6 +51,7 @@ func TestStorage_New(t *testing.T) {
 				assert.Equal(s.(*storage).bufferSize, DefaultBufferSize)
 				assert.Equal(cap(s.(*storage).buffer), DefaultBufferSize)
 				assert.Equal(len(s.(*storage).buffer), 0)
+				assert.Equal(s.(*storage).count, int64(0))
 
 				if err := s.Clear(); err != nil {
 					t.Fatal(err)
@@ -108,6 +109,7 @@ func TestStorage_New(t *testing.T) {
 				assert.Equal(s.(*storage).bufferSize, 1)
 				assert.Equal(cap(s.(*storage).buffer), 1)
 				assert.Equal(len(s.(*storage).buffer), 0)
+				assert.Equal(len(s.(*storage).buffer), 0)
 
 				if err := s.Clear(); err != nil {
 					t.Fatal(err)
@@ -150,6 +152,35 @@ func TestStorage_Create(t *testing.T) {
 				assert := assert.New(t)
 				err := s.Create(Record{})
 				assert.NoError(err)
+				assert.Equal(s.(*storage).count, int64(0))
+			},
+		},
+		{
+			name:    "create record without buffer",
+			baseDir: os.TempDir(),
+			options: []Option{WithBufferSize(0)},
+			mock: func(s Storage) {
+			},
+			expect: func(t *testing.T, s Storage, baseDir string) {
+				assert := assert.New(t)
+				err := s.Create(Record{})
+				assert.NoError(err)
+				assert.Equal(s.(*storage).count, int64(1))
+			},
+		},
+		{
+			name:    "write record to file",
+			baseDir: os.TempDir(),
+			options: []Option{WithBufferSize(1)},
+			mock: func(s Storage) {
+			},
+			expect: func(t *testing.T, s Storage, baseDir string) {
+				assert := assert.New(t)
+				err := s.Create(Record{})
+				assert.NoError(err)
+				err = s.Create(Record{})
+				assert.NoError(err)
+				assert.Equal(s.(*storage).count, int64(1))
 			},
 		},
 		{