Add Zygostates hardware qualification runner
This commit is contained in:
commit
65541dc8eb
1 changed files with 2889 additions and 0 deletions
157
internal/zygostates/uploader.go
Normal file
157
internal/zygostates/uploader.go
Normal file
|
|
@ -0,0 +1,157 @@
|
|||
package zygostates
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"compress/gzip"
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/url"
|
||||
"path"
|
||||
"strconv"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
type PutRequest struct {
|
||||
Bucket string
|
||||
Key string
|
||||
Body []byte
|
||||
ContentType string
|
||||
ContentEncoding string
|
||||
Metadata map[string]string
|
||||
}
|
||||
|
||||
type ObjectStore interface {
|
||||
Put(context.Context, PutRequest) error
|
||||
}
|
||||
|
||||
type UploadedRange struct {
|
||||
Path string `json:"path"`
|
||||
Generation int `json:"generation"`
|
||||
Start int64 `json:"start"`
|
||||
End int64 `json:"end"`
|
||||
SHA256 string `json:"sha256"`
|
||||
Key string `json:"key"`
|
||||
}
|
||||
|
||||
type UploadTick struct {
|
||||
Timestamp time.Time `json:"timestamp"`
|
||||
Sequence uint64 `json:"sequence"`
|
||||
Uploaded []UploadedRange `json:"uploaded"`
|
||||
}
|
||||
|
||||
type uploadPosition struct {
|
||||
offset int64
|
||||
generation int
|
||||
}
|
||||
|
||||
type Uploader struct {
|
||||
store ObjectStore
|
||||
artifacts *ArtifactStore
|
||||
bucket string
|
||||
baseKey string
|
||||
positions map[string]uploadPosition
|
||||
sequence uint64
|
||||
maxChunk int
|
||||
mu sync.Mutex
|
||||
}
|
||||
|
||||
func newUploader(store ObjectStore, artifacts *ArtifactStore, config UploadConfig, identity, runID string) *Uploader {
|
||||
return &Uploader{
|
||||
store: store, artifacts: artifacts, bucket: config.Bucket,
|
||||
baseKey: path.Join(config.Prefix, identity, runID),
|
||||
positions: make(map[string]uploadPosition), maxChunk: 8 * 1024 * 1024,
|
||||
}
|
||||
}
|
||||
|
||||
func (u *Uploader) Upload(ctx context.Context) (int, error) {
|
||||
u.mu.Lock()
|
||||
defer u.mu.Unlock()
|
||||
snapshots, err := u.artifacts.Snapshots()
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
staged := make(map[string]uploadPosition, len(u.positions))
|
||||
for name, position := range u.positions {
|
||||
staged[name] = position
|
||||
}
|
||||
var uploaded []UploadedRange
|
||||
for _, snapshot := range snapshots {
|
||||
position := staged[snapshot.Name]
|
||||
if int64(len(snapshot.Contents)) < position.offset {
|
||||
position.offset = 0
|
||||
position.generation++
|
||||
}
|
||||
for position.offset < int64(len(snapshot.Contents)) {
|
||||
end := min(position.offset+int64(u.maxChunk), int64(len(snapshot.Contents)))
|
||||
content := snapshot.Contents[position.offset:end]
|
||||
digest := sha256.Sum256(content)
|
||||
digestText := hex.EncodeToString(digest[:])
|
||||
key := path.Join(u.baseKey, "files", url.PathEscape(snapshot.Name), fmt.Sprintf("%04d", position.generation), fmt.Sprintf("%020d-%020d.gz", position.offset, end))
|
||||
var compressed bytes.Buffer
|
||||
writer, _ := gzip.NewWriterLevel(&compressed, gzip.BestSpeed)
|
||||
if _, err := writer.Write(content); err != nil {
|
||||
return len(uploaded), err
|
||||
}
|
||||
if err := writer.Close(); err != nil {
|
||||
return len(uploaded), err
|
||||
}
|
||||
if err := u.store.Put(ctx, PutRequest{
|
||||
Bucket: u.bucket, Key: key, Body: compressed.Bytes(), ContentType: "application/octet-stream", ContentEncoding: "gzip",
|
||||
Metadata: map[string]string{"source-path": snapshot.Name, "generation": strconv.Itoa(position.generation), "start": strconv.FormatInt(position.offset, 10), "end": strconv.FormatInt(end, 10), "sha256": digestText},
|
||||
}); err != nil {
|
||||
return len(uploaded), err
|
||||
}
|
||||
uploaded = append(uploaded, UploadedRange{Path: snapshot.Name, Generation: position.generation, Start: position.offset, End: end, SHA256: digestText, Key: key})
|
||||
position.offset = end
|
||||
}
|
||||
staged[snapshot.Name] = position
|
||||
}
|
||||
tick := UploadTick{Timestamp: time.Now().UTC(), Sequence: u.sequence, Uploaded: uploaded}
|
||||
tickBody, err := json.Marshal(tick)
|
||||
if err != nil {
|
||||
return len(uploaded), err
|
||||
}
|
||||
tickKey := path.Join(u.baseKey, "ticks", fmt.Sprintf("%08d.json", u.sequence))
|
||||
if err := u.store.Put(ctx, PutRequest{Bucket: u.bucket, Key: tickKey, Body: tickBody, ContentType: "application/json"}); err != nil {
|
||||
return len(uploaded), err
|
||||
}
|
||||
u.positions = staged
|
||||
u.sequence++
|
||||
return len(uploaded), nil
|
||||
}
|
||||
|
||||
type PeriodicUploader struct {
|
||||
uploader *Uploader
|
||||
interval time.Duration
|
||||
store *ArtifactStore
|
||||
}
|
||||
|
||||
func (u *PeriodicUploader) Start(ctx context.Context) *sync.WaitGroup {
|
||||
var group sync.WaitGroup
|
||||
group.Add(1)
|
||||
go func() {
|
||||
defer group.Done()
|
||||
ticker := time.NewTicker(u.interval)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
uploadContext, cancel := context.WithTimeout(context.WithoutCancel(ctx), 30*time.Second)
|
||||
count, err := u.uploader.Upload(uploadContext)
|
||||
cancel()
|
||||
fields := map[string]any{"uploaded_ranges": count}
|
||||
if err != nil {
|
||||
fields["error"] = err.Error()
|
||||
}
|
||||
_ = u.store.AppendEvent(Event{Timestamp: time.Now().UTC(), Type: "upload", Fields: fields})
|
||||
}
|
||||
}
|
||||
}()
|
||||
return &group
|
||||
}
|
||||
Loading…
Reference in a new issue