Dragonfly2

1
/*
2
 *     Copyright 2023 The Dragonfly Authors
3
 *
4
 * Licensed under the Apache License, Version 2.0 (the "License");
5
 * you may not use this file except in compliance with the License.
6
 * You may obtain a copy of the License at
7
 *
8
 *      http://www.apache.org/licenses/LICENSE-2.0
9
 *
10
 * Unless required by applicable law or agreed to in writing, software
11
 * distributed under the License is distributed on an "AS IS" BASIS,
12
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
 * See the License for the specific language governing permissions and
14
 * limitations under the License.
15
 */
16

17
//go:generate mockgen -destination mocks/storage_mock.go -source storage.go -package mocks
18

19
package storage
20

21
import (
22
	"errors"
23
	"fmt"
24
	"os"
25
	"path/filepath"
26

27
	"github.com/gocarina/gocsv"
28

29
	schedulerstorage "d7y.io/dragonfly/v2/scheduler/storage"
30
)
31

32
const (
33
	// DownloadFilePrefix is prefix of download file name.
34
	DownloadFilePrefix = "download"
35

36
	// NetworkTopologyFilePrefix is prefix of network topology file name.
37
	NetworkTopologyFilePrefix = "networktopology"
38

39
	// CSVFileExt is extension of file name.
40
	CSVFileExt = "csv"
41
)
42

43
// Storage is the interface used for storage.
44
type Storage interface {
45
	// ListDownload returns downloads in csv files based on the given model key.
46
	ListDownload(string) ([]schedulerstorage.Download, error)
47

48
	// ListNetworkTopology returns network topologies in csv files based on the given model key.
49
	ListNetworkTopology(string) ([]schedulerstorage.NetworkTopology, error)
50

51
	// OpenDownload opens download files for read based on the given model key, it returns io.ReadCloser of download files.
52
	OpenDownload(string) (*os.File, error)
53

54
	// OpenNetworkTopology opens network topology files for read based on the given model key, it returns io.ReadCloser of network topology files.
55
	OpenNetworkTopology(string) (*os.File, error)
56

57
	// ClearDownload removes all downloads based on the given model key.
58
	ClearDownload(string) error
59

60
	// ClearNetworkTopology removes network topologies based on the given model key.
61
	ClearNetworkTopology(string) error
62

63
	// Clear removes all files.
64
	Clear() error
65
}
66

67
// storage provides storage function.
68
type storage struct {
69
	baseDir string
70
}
71

72
// New returns a new Storage instance.
73
func New(baseDir string) Storage {
74
	return &storage{baseDir: baseDir}
75
}
76

77
// ListDownload returns downloads in csv files based on the given model key.
78
func (s *storage) ListDownload(key string) (downloads []schedulerstorage.Download, err error) {
79
	file, err := s.OpenDownload(key)
80
	if err != nil {
81
		return nil, err
82
	}
83
	defer func() {
84
		if cerr := file.Close(); cerr != nil {
85
			err = errors.Join(err, cerr)
86
		}
87
	}()
88

89
	if err = gocsv.UnmarshalWithoutHeaders(file, &downloads); err != nil {
90
		return nil, err
91
	}
92

93
	return downloads, nil
94
}
95

96
// ListNetworkTopology returns network topologies in csv files based on the given model key.
97
func (s *storage) ListNetworkTopology(key string) (networkTopologies []schedulerstorage.NetworkTopology, err error) {
98
	file, err := s.OpenNetworkTopology(key)
99
	if err != nil {
100
		return nil, err
101
	}
102
	defer func() {
103
		if cerr := file.Close(); cerr != nil {
104
			err = errors.Join(err, cerr)
105
		}
106
	}()
107

108
	if err = gocsv.UnmarshalWithoutHeaders(file, &networkTopologies); err != nil {
109
		return nil, err
110
	}
111

112
	return networkTopologies, nil
113
}
114

115
// OpenDownload opens download files for read based on the given model key, it returns io.ReadCloser of download files.
116
func (s *storage) OpenDownload(key string) (*os.File, error) {
117
	return os.OpenFile(s.downloadFilename(key), os.O_RDWR|os.O_CREATE|os.O_APPEND, 0600)
118
}
119

120
// OpenNetworkTopology opens network topology files for read based on the given model key, it returns io.ReadCloser of network topology files.
121
func (s *storage) OpenNetworkTopology(key string) (*os.File, error) {
122
	return os.OpenFile(s.networkTopologyFilename(key), os.O_RDWR|os.O_CREATE|os.O_APPEND, 0600)
123
}
124

125
// ClearDownload removes downloads based on the given model key.
126
func (s *storage) ClearDownload(key string) error {
127
	return os.Remove(s.downloadFilename(key))
128
}
129

130
// ClearNetworkTopology removes network topologies based on the given model key.
131
func (s *storage) ClearNetworkTopology(key string) error {
132
	return os.Remove(s.networkTopologyFilename(key))
133
}
134

135
// Clear removes all files.
136
func (s *storage) Clear() error {
137
	return os.RemoveAll(s.baseDir)
138
}
139

140
// downloadFilename generates download file name based on the given model key.
141
func (s *storage) downloadFilename(key string) string {
142
	return filepath.Join(s.baseDir, fmt.Sprintf("%s_%s.%s", DownloadFilePrefix, key, CSVFileExt))
143
}
144

145
// networkTopologyFilename generates network topology file name based on the given model key.
146
func (s *storage) networkTopologyFilename(key string) string {
147
	return filepath.Join(s.baseDir, fmt.Sprintf("%s_%s.%s", NetworkTopologyFilePrefix, key, CSVFileExt))
148
}
149