kraken

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

16
import (
17
	"fmt"
18
	"regexp"
19
	"sync"
20

21
	"github.com/uber/kraken/core"
22
	"github.com/uber/kraken/lib/store/metadata"
23
	"github.com/uber/kraken/utils/log"
24
)
25

26
const _pieceStatusSuffix = "_status"
27

28
func init() {
29
	metadata.Register(regexp.MustCompile(_pieceStatusSuffix), pieceStatusMetadataFactory{})
30
}
31

32
type pieceStatus int
33

34
const (
35
	_empty pieceStatus = iota
36
	_complete
37
	_dirty
38
)
39

40
type pieceStatusMetadataFactory struct{}
41

42
func (m pieceStatusMetadataFactory) Create(suffix string) metadata.Metadata {
43
	return &pieceStatusMetadata{}
44
}
45

46
// pieceStatusMetadata stores pieces statuses as metadata on disk.
47
type pieceStatusMetadata struct {
48
	pieces []*piece
49
}
50

51
func newPieceStatusMetadata(pieces []*piece) *pieceStatusMetadata {
52
	return &pieceStatusMetadata{pieces}
53
}
54

55
func (m *pieceStatusMetadata) GetSuffix() string {
56
	return _pieceStatusSuffix
57
}
58

59
func (m *pieceStatusMetadata) Movable() bool {
60
	return true
61
}
62

63
func (m *pieceStatusMetadata) Serialize() ([]byte, error) {
64
	b := make([]byte, len(m.pieces))
65
	for i, p := range m.pieces {
66
		b[i] = byte(p.status)
67
	}
68
	return b, nil
69
}
70

71
func (m *pieceStatusMetadata) Deserialize(b []byte) error {
72
	m.pieces = make([]*piece, len(b))
73
	for i := range b {
74
		status := pieceStatus(b[i])
75
		if status != _empty && status != _complete {
76
			log.Errorf("Unexpected status in piece metadata: %d", status)
77
			status = _empty
78
		}
79
		m.pieces[i] = &piece{status: status}
80
	}
81
	return nil
82
}
83

84
type piece struct {
85
	sync.RWMutex
86
	status pieceStatus
87
}
88

89
func (p *piece) complete() bool {
90
	p.RLock()
91
	defer p.RUnlock()
92
	return p.status == _complete
93
}
94

95
func (p *piece) dirty() bool {
96
	p.RLock()
97
	defer p.RUnlock()
98
	return p.status == _dirty
99
}
100

101
func (p *piece) tryMarkDirty() (dirty, complete bool) {
102
	p.Lock()
103
	defer p.Unlock()
104

105
	switch p.status {
106
	case _empty:
107
		p.status = _dirty
108
	case _dirty:
109
		dirty = true
110
	case _complete:
111
		complete = true
112
	default:
113
		log.Fatalf("Unknown piece status: %d", p.status)
114
	}
115
	return
116
}
117

118
func (p *piece) markEmpty() {
119
	p.Lock()
120
	defer p.Unlock()
121
	p.status = _empty
122
}
123

124
func (p *piece) markComplete() {
125
	p.Lock()
126
	defer p.Unlock()
127
	p.status = _complete
128
}
129

130
// restorePieces reads piece metadata from disk and restores the in-memory piece
131
// statuses. A naive solution would be to read the entire blob from disk and
132
// hash the pieces to determine completion status -- however, this is very
133
// expensive. Instead, Torrent tracks completed pieces on disk via metadata
134
// as they are written.
135
func restorePieces(
136
	d core.Digest,
137
	cads caDownloadStore,
138
	numPieces int) (pieces []*piece, numComplete int, err error) {
139

140
	for i := 0; i < numPieces; i++ {
141
		pieces = append(pieces, &piece{status: _empty})
142
	}
143
	md := newPieceStatusMetadata(pieces)
144
	if err := cads.Download().GetOrSetMetadata(d.Hex(), md); cads.InCacheError(err) {
145
		// File is in cache state -- initialize completed pieces.
146
		for _, p := range pieces {
147
			p.status = _complete
148
		}
149
		return pieces, numPieces, nil
150
	} else if err != nil {
151
		return nil, 0, fmt.Errorf("get or set piece metadata: %s", err)
152
	}
153
	for _, p := range md.pieces {
154
		if p.status == _complete {
155
			numComplete++
156
		}
157
	}
158
	return md.pieces, numComplete, nil
159
}
160