-
Notifications
You must be signed in to change notification settings - Fork 1
/
main.go
235 lines (204 loc) · 6.28 KB
/
main.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
// The resource migrator allow us to migrate all the resources saved on disk to a storage provider such as S3.
package main
import (
"context"
"github.com/aws/aws-sdk-go/aws/session"
"github.com/aws/aws-sdk-go/service/s3"
"github.com/aws/aws-sdk-go/service/s3/s3manager"
res "github.com/gazebo-web/fuel-server/bundles/common_resources"
"github.com/gazebo-web/fuel-server/bundles/models"
"github.com/gazebo-web/fuel-server/bundles/worlds"
"github.com/gazebo-web/fuel-server/globals"
"github.com/gazebo-web/fuel-server/vcs"
"github.com/gazebo-web/gz-go/v7"
"github.com/gazebo-web/gz-go/v7/storage"
"github.com/jinzhu/gorm"
_ "github.com/jinzhu/gorm/dialects/mysql"
"github.com/schollz/progressbar/v3"
"log"
"os"
"os/signal"
"strconv"
"syscall"
"time"
)
const (
maxParallelUploads = 15
)
func main() {
s, db := setup()
defer gz.Close(db)
run(s, db)
}
func run(s storage.Storage, db *gorm.DB) {
started := time.Now()
// Get the list of resources to upload for: Models
ms, modelListSize, err := getResourceToUpload[*models.Model](db)
if err != nil {
log.Panicln("Failed to get models:", err)
}
// Get the list of resources to upload for: Worlds
ws, worldListSize, err := getResourceToUpload[*worlds.World](db)
if err != nil {
log.Panicln("Failed to get worlds:", err)
}
totalSize := modelListSize + worldListSize
// Channel to request resource uploads
c := make(chan uploadRequest, totalSize)
// Channel to handle errors
e := make(chan uploadError, totalSize)
// Requesting all models
log.Println("Processing Models")
requestUpload[*models.Model](c, ms, "models")
// Requesting all worlds
log.Println("Processing Worlds")
requestUpload[*worlds.World](c, ws, "worlds")
// Listen for exit signal from when all resources have been uploaded
exit := make(chan struct{}, 1)
// Begin parallel uploads and keep track of them using the progress bar. It will send a single item to the exit
// channel once it finished
bar := newProgressBar(totalSize, exit)
for i := 0; i < maxParallelUploads; i++ {
go upload(c, e, s, bar)
}
// Listen for Interrupt and Terminate signals
sigs := make(chan os.Signal, 1)
signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)
fails := make([]uploadError, 0, modelListSize+worldListSize)
for {
select {
case failed := <-e:
fails = append(fails, failed)
case sig := <-sigs:
log.Panicln("Signal received:", sig.String())
case <-exit:
log.Println("Models and Worlds were migrated. Took:", time.Since(started).Seconds(), "seconds")
if len(fails) > 0 {
log.Printf("However, the following resources returned an error while uploading them:")
for _, fail := range fails {
log.Printf("Resource [%s]: %s/%s - Error: %s\n", fail.Request.Kind,
*fail.Request.Resource.GetOwner(), *fail.Request.Resource.GetName(), fail.Error)
}
}
return
}
}
}
func setup() (storage.Storage, *gorm.DB) {
db, err := setupDB()
if err != nil {
log.Fatalln("Failed to set up to MySQL database conn:", err)
}
globals.ResourceDir = os.Getenv("FUEL_RESOURCE_DIR")
// Set up git
globals.VCSRepoFactory = func(ctx context.Context, dirpath string) vcs.VCS {
return vcs.GoGitVCS{}.NewRepo(dirpath)
}
// Initialize S3 config
s3session := session.Must(session.NewSession())
s := storage.NewS3v1(s3.New(s3session), s3manager.NewUploader(s3session), os.Getenv("AWS_S3_BUCKET"))
return s, db
}
func newProgressBar(size int, exit chan struct{}) *progressbar.ProgressBar {
return progressbar.NewOptions(size,
progressbar.OptionSetPredictTime(false),
progressbar.OptionShowCount(),
progressbar.OptionSetDescription("Uploading resources"),
progressbar.OptionSetWriter(os.Stderr),
progressbar.OptionSetWidth(10),
progressbar.OptionThrottle(65*time.Millisecond),
progressbar.OptionOnCompletion(func() {
exit <- struct{}{}
}),
progressbar.OptionSpinnerType(14),
progressbar.OptionFullWidth(),
progressbar.OptionSetRenderBlankState(true),
)
}
func setupDB() (*gorm.DB, error) {
// Initialize database
cfg, err := gz.NewDatabaseConfigFromEnvVars()
if err != nil {
return nil, err
}
db, err := gz.InitDbWithCfg(&cfg)
if err != nil {
return nil, err
}
return db, err
}
type uploadRequest struct {
Kind string
Resource res.Resource
}
type uploadError struct {
Error error
Request uploadRequest
}
func requestUpload[T res.Resource](c chan uploadRequest, items []T, kind string) {
for _, item := range items {
c <- uploadRequest{
Kind: kind,
Resource: item,
}
}
}
func upload(c chan uploadRequest, e chan uploadError, storage storage.Storage, bar *progressbar.ProgressBar) {
for !bar.IsFinished() {
req := <-c
err := uploadResources(context.Background(), storage, req.Kind, req.Resource)
if err != nil {
e <- uploadError{
Error: err,
Request: req,
}
}
_ = bar.Add(1)
}
}
func getResourceToUpload[T res.Resource](db *gorm.DB) ([]T, int, error) {
var list []T
var model T
if err := db.Model(&model).Find(&list).Error; err != nil {
return nil, 0, err
}
return list, len(list), nil
}
func uploadResources(ctx context.Context, storage storage.Storage, kind string, r res.Resource) error {
v, err := uploadResource(ctx, storage, kind, "tip", r)
if err != nil {
return err
}
// If the tip version is not the version 1, we should migrate all the older versions.
for v > 1 {
// Decrease by 1
v--
// Upload the resources for the current versions
v, err = uploadResource(ctx, storage, kind, strconv.Itoa(v), r)
if err != nil {
return err
}
}
return nil
}
func uploadResource(ctx context.Context, storage storage.Storage, kind, versions string, r res.Resource) (int, error) {
path, ver, em := res.GetZip(ctx, r, kind, versions)
if em != nil {
log.Printf("Failed to get zip file for %s: %s\n", kind, em.BaseError)
return 0, em.BaseError
}
f, err := os.Open(*path)
defer gz.Close(f)
if err != nil {
log.Printf("Failed to open zip file for %s: %s\n", kind, err)
log.Printf("Name: %s | Owner: %s | versions: %d | Path: %d\n", *r.GetName(), *r.GetOwner(), ver, path)
return 0, err
}
err = storage.UploadZip(ctx, res.CastResourceToStorageResource(r, uint64(ver)), f)
if err != nil {
log.Printf("Failed to upload zip file for %s: %s\n", kind, err)
log.Printf("Name: %s | Owner: %s | versions: %d | Path: %d\n", *r.GetName(), *r.GetOwner(), ver, path)
return 0, err
}
return ver, nil
}