github.com/walkingsparrow/docker@v1.4.2-0.20151218153551-b708a2249bfa/distribution/xfer/upload_test.go (about)

     1  package xfer
     2  
     3  import (
     4  	"errors"
     5  	"sync/atomic"
     6  	"testing"
     7  	"time"
     8  
     9  	"github.com/docker/distribution/digest"
    10  	"github.com/docker/docker/layer"
    11  	"github.com/docker/docker/pkg/progress"
    12  	"golang.org/x/net/context"
    13  )
    14  
    15  const maxUploadConcurrency = 3
    16  
    17  type mockUploadDescriptor struct {
    18  	currentUploads  *int32
    19  	diffID          layer.DiffID
    20  	simulateRetries int
    21  }
    22  
    23  // Key returns the key used to deduplicate downloads.
    24  func (u *mockUploadDescriptor) Key() string {
    25  	return u.diffID.String()
    26  }
    27  
    28  // ID returns the ID for display purposes.
    29  func (u *mockUploadDescriptor) ID() string {
    30  	return u.diffID.String()
    31  }
    32  
    33  // DiffID should return the DiffID for this layer.
    34  func (u *mockUploadDescriptor) DiffID() layer.DiffID {
    35  	return u.diffID
    36  }
    37  
    38  // Upload is called to perform the upload.
    39  func (u *mockUploadDescriptor) Upload(ctx context.Context, progressOutput progress.Output) (digest.Digest, error) {
    40  	if u.currentUploads != nil {
    41  		defer atomic.AddInt32(u.currentUploads, -1)
    42  
    43  		if atomic.AddInt32(u.currentUploads, 1) > maxUploadConcurrency {
    44  			return "", errors.New("concurrency limit exceeded")
    45  		}
    46  	}
    47  
    48  	// Sleep a bit to simulate a time-consuming upload.
    49  	for i := int64(0); i <= 10; i++ {
    50  		select {
    51  		case <-ctx.Done():
    52  			return "", ctx.Err()
    53  		case <-time.After(10 * time.Millisecond):
    54  			progressOutput.WriteProgress(progress.Progress{ID: u.ID(), Current: i, Total: 10})
    55  		}
    56  	}
    57  
    58  	if u.simulateRetries != 0 {
    59  		u.simulateRetries--
    60  		return "", errors.New("simulating retry")
    61  	}
    62  
    63  	// For the mock implementation, use SHA256(DiffID) as the returned
    64  	// digest.
    65  	return digest.FromBytes([]byte(u.diffID.String()))
    66  }
    67  
    68  func uploadDescriptors(currentUploads *int32) []UploadDescriptor {
    69  	return []UploadDescriptor{
    70  		&mockUploadDescriptor{currentUploads, layer.DiffID("sha256:cbbf2f9a99b47fc460d422812b6a5adff7dfee951d8fa2e4a98caa0382cfbdbf"), 0},
    71  		&mockUploadDescriptor{currentUploads, layer.DiffID("sha256:1515325234325236634634608943609283523908626098235490238423902343"), 0},
    72  		&mockUploadDescriptor{currentUploads, layer.DiffID("sha256:6929356290463485374960346430698374523437683470934634534953453453"), 0},
    73  		&mockUploadDescriptor{currentUploads, layer.DiffID("sha256:cbbf2f9a99b47fc460d422812b6a5adff7dfee951d8fa2e4a98caa0382cfbdbf"), 0},
    74  		&mockUploadDescriptor{currentUploads, layer.DiffID("sha256:8159352387436803946235346346368745389534789534897538734598734987"), 1},
    75  		&mockUploadDescriptor{currentUploads, layer.DiffID("sha256:4637863963478346897346987346987346789346789364879364897364987346"), 0},
    76  	}
    77  }
    78  
    79  var expectedDigests = map[layer.DiffID]digest.Digest{
    80  	layer.DiffID("sha256:cbbf2f9a99b47fc460d422812b6a5adff7dfee951d8fa2e4a98caa0382cfbdbf"): digest.Digest("sha256:c5095d6cf7ee42b7b064371dcc1dc3fb4af197f04d01a60009d484bd432724fc"),
    81  	layer.DiffID("sha256:1515325234325236634634608943609283523908626098235490238423902343"): digest.Digest("sha256:968cbfe2ff5269ea1729b3804767a1f57ffbc442d3bc86f47edbf7e688a4f36e"),
    82  	layer.DiffID("sha256:6929356290463485374960346430698374523437683470934634534953453453"): digest.Digest("sha256:8a5e56ab4b477a400470a7d5d4c1ca0c91235fd723ab19cc862636a06f3a735d"),
    83  	layer.DiffID("sha256:8159352387436803946235346346368745389534789534897538734598734987"): digest.Digest("sha256:5e733e5cd3688512fc240bd5c178e72671c9915947d17bb8451750d827944cb2"),
    84  	layer.DiffID("sha256:4637863963478346897346987346987346789346789364879364897364987346"): digest.Digest("sha256:ec4bb98d15e554a9f66c3ef9296cf46772c0ded3b1592bd8324d96e2f60f460c"),
    85  }
    86  
    87  func TestSuccessfulUpload(t *testing.T) {
    88  	lum := NewLayerUploadManager(maxUploadConcurrency)
    89  
    90  	progressChan := make(chan progress.Progress)
    91  	progressDone := make(chan struct{})
    92  	receivedProgress := make(map[string]int64)
    93  
    94  	go func() {
    95  		for p := range progressChan {
    96  			receivedProgress[p.ID] = p.Current
    97  		}
    98  		close(progressDone)
    99  	}()
   100  
   101  	var currentUploads int32
   102  	descriptors := uploadDescriptors(&currentUploads)
   103  
   104  	digests, err := lum.Upload(context.Background(), descriptors, progress.ChanOutput(progressChan))
   105  	if err != nil {
   106  		t.Fatalf("upload error: %v", err)
   107  	}
   108  
   109  	close(progressChan)
   110  	<-progressDone
   111  
   112  	if len(digests) != len(expectedDigests) {
   113  		t.Fatal("wrong number of keys in digests map")
   114  	}
   115  
   116  	for key, val := range expectedDigests {
   117  		if digests[key] != val {
   118  			t.Fatalf("mismatch in digest array for key %v (expected %v, got %v)", key, val, digests[key])
   119  		}
   120  		if receivedProgress[key.String()] != 10 {
   121  			t.Fatalf("missing or wrong progress output for %v", key)
   122  		}
   123  	}
   124  }
   125  
   126  func TestCancelledUpload(t *testing.T) {
   127  	lum := NewLayerUploadManager(maxUploadConcurrency)
   128  
   129  	progressChan := make(chan progress.Progress)
   130  	progressDone := make(chan struct{})
   131  
   132  	go func() {
   133  		for range progressChan {
   134  		}
   135  		close(progressDone)
   136  	}()
   137  
   138  	ctx, cancel := context.WithCancel(context.Background())
   139  
   140  	go func() {
   141  		<-time.After(time.Millisecond)
   142  		cancel()
   143  	}()
   144  
   145  	descriptors := uploadDescriptors(nil)
   146  	_, err := lum.Upload(ctx, descriptors, progress.ChanOutput(progressChan))
   147  	if err != context.Canceled {
   148  		t.Fatal("expected upload to be cancelled")
   149  	}
   150  
   151  	close(progressChan)
   152  	<-progressDone
   153  }