github.com/uber/kraken@v0.1.4/lib/torrent/storage/agentstorage/pieces.go (about)

     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  }