github.com/rita33cool1/iot-system-gateway@v0.0.0-20200911033302-e65bde238cc5/docker-engine/distribution/xfer/download_test.go (about)

     1  package xfer // import "github.com/docker/docker/distribution/xfer"
     2  
     3  import (
     4  	"bytes"
     5  	"errors"
     6  	"fmt"
     7  	"io"
     8  	"io/ioutil"
     9  	"runtime"
    10  	"sync/atomic"
    11  	"testing"
    12  	"time"
    13  
    14  	"github.com/docker/distribution"
    15  	"github.com/docker/docker/image"
    16  	"github.com/docker/docker/layer"
    17  	"github.com/docker/docker/pkg/progress"
    18  	"github.com/opencontainers/go-digest"
    19  	"golang.org/x/net/context"
    20  )
    21  
    22  const maxDownloadConcurrency = 3
    23  
    24  type mockLayer struct {
    25  	layerData bytes.Buffer
    26  	diffID    layer.DiffID
    27  	chainID   layer.ChainID
    28  	parent    layer.Layer
    29  	os        string
    30  }
    31  
    32  func (ml *mockLayer) TarStream() (io.ReadCloser, error) {
    33  	return ioutil.NopCloser(bytes.NewBuffer(ml.layerData.Bytes())), nil
    34  }
    35  
    36  func (ml *mockLayer) TarStreamFrom(layer.ChainID) (io.ReadCloser, error) {
    37  	return nil, fmt.Errorf("not implemented")
    38  }
    39  
    40  func (ml *mockLayer) ChainID() layer.ChainID {
    41  	return ml.chainID
    42  }
    43  
    44  func (ml *mockLayer) DiffID() layer.DiffID {
    45  	return ml.diffID
    46  }
    47  
    48  func (ml *mockLayer) Parent() layer.Layer {
    49  	return ml.parent
    50  }
    51  
    52  func (ml *mockLayer) Size() (size int64, err error) {
    53  	return 0, nil
    54  }
    55  
    56  func (ml *mockLayer) DiffSize() (size int64, err error) {
    57  	return 0, nil
    58  }
    59  
    60  func (ml *mockLayer) Metadata() (map[string]string, error) {
    61  	return make(map[string]string), nil
    62  }
    63  
    64  type mockLayerStore struct {
    65  	layers map[layer.ChainID]*mockLayer
    66  }
    67  
    68  func createChainIDFromParent(parent layer.ChainID, dgsts ...layer.DiffID) layer.ChainID {
    69  	if len(dgsts) == 0 {
    70  		return parent
    71  	}
    72  	if parent == "" {
    73  		return createChainIDFromParent(layer.ChainID(dgsts[0]), dgsts[1:]...)
    74  	}
    75  	// H = "H(n-1) SHA256(n)"
    76  	dgst := digest.FromBytes([]byte(string(parent) + " " + string(dgsts[0])))
    77  	return createChainIDFromParent(layer.ChainID(dgst), dgsts[1:]...)
    78  }
    79  
    80  func (ls *mockLayerStore) Map() map[layer.ChainID]layer.Layer {
    81  	layers := map[layer.ChainID]layer.Layer{}
    82  
    83  	for k, v := range ls.layers {
    84  		layers[k] = v
    85  	}
    86  
    87  	return layers
    88  }
    89  
    90  func (ls *mockLayerStore) Register(reader io.Reader, parentID layer.ChainID) (layer.Layer, error) {
    91  	return ls.RegisterWithDescriptor(reader, parentID, distribution.Descriptor{})
    92  }
    93  
    94  func (ls *mockLayerStore) RegisterWithDescriptor(reader io.Reader, parentID layer.ChainID, _ distribution.Descriptor) (layer.Layer, error) {
    95  	var (
    96  		parent layer.Layer
    97  		err    error
    98  	)
    99  
   100  	if parentID != "" {
   101  		parent, err = ls.Get(parentID)
   102  		if err != nil {
   103  			return nil, err
   104  		}
   105  	}
   106  
   107  	l := &mockLayer{parent: parent}
   108  	_, err = l.layerData.ReadFrom(reader)
   109  	if err != nil {
   110  		return nil, err
   111  	}
   112  	l.diffID = layer.DiffID(digest.FromBytes(l.layerData.Bytes()))
   113  	l.chainID = createChainIDFromParent(parentID, l.diffID)
   114  
   115  	ls.layers[l.chainID] = l
   116  	return l, nil
   117  }
   118  
   119  func (ls *mockLayerStore) Get(chainID layer.ChainID) (layer.Layer, error) {
   120  	l, ok := ls.layers[chainID]
   121  	if !ok {
   122  		return nil, layer.ErrLayerDoesNotExist
   123  	}
   124  	return l, nil
   125  }
   126  
   127  func (ls *mockLayerStore) Release(l layer.Layer) ([]layer.Metadata, error) {
   128  	return []layer.Metadata{}, nil
   129  }
   130  func (ls *mockLayerStore) CreateRWLayer(string, layer.ChainID, *layer.CreateRWLayerOpts) (layer.RWLayer, error) {
   131  	return nil, errors.New("not implemented")
   132  }
   133  
   134  func (ls *mockLayerStore) GetRWLayer(string) (layer.RWLayer, error) {
   135  	return nil, errors.New("not implemented")
   136  }
   137  
   138  func (ls *mockLayerStore) ReleaseRWLayer(layer.RWLayer) ([]layer.Metadata, error) {
   139  	return nil, errors.New("not implemented")
   140  }
   141  func (ls *mockLayerStore) GetMountID(string) (string, error) {
   142  	return "", errors.New("not implemented")
   143  }
   144  
   145  func (ls *mockLayerStore) Cleanup() error {
   146  	return nil
   147  }
   148  
   149  func (ls *mockLayerStore) DriverStatus() [][2]string {
   150  	return [][2]string{}
   151  }
   152  
   153  func (ls *mockLayerStore) DriverName() string {
   154  	return "mock"
   155  }
   156  
   157  type mockDownloadDescriptor struct {
   158  	currentDownloads *int32
   159  	id               string
   160  	diffID           layer.DiffID
   161  	registeredDiffID layer.DiffID
   162  	expectedDiffID   layer.DiffID
   163  	simulateRetries  int
   164  }
   165  
   166  // Key returns the key used to deduplicate downloads.
   167  func (d *mockDownloadDescriptor) Key() string {
   168  	return d.id
   169  }
   170  
   171  // ID returns the ID for display purposes.
   172  func (d *mockDownloadDescriptor) ID() string {
   173  	return d.id
   174  }
   175  
   176  // DiffID should return the DiffID for this layer, or an error
   177  // if it is unknown (for example, if it has not been downloaded
   178  // before).
   179  func (d *mockDownloadDescriptor) DiffID() (layer.DiffID, error) {
   180  	if d.diffID != "" {
   181  		return d.diffID, nil
   182  	}
   183  	return "", errors.New("no diffID available")
   184  }
   185  
   186  func (d *mockDownloadDescriptor) Registered(diffID layer.DiffID) {
   187  	d.registeredDiffID = diffID
   188  }
   189  
   190  func (d *mockDownloadDescriptor) mockTarStream() io.ReadCloser {
   191  	// The mock implementation returns the ID repeated 5 times as a tar
   192  	// stream instead of actual tar data. The data is ignored except for
   193  	// computing IDs.
   194  	return ioutil.NopCloser(bytes.NewBuffer([]byte(d.id + d.id + d.id + d.id + d.id)))
   195  }
   196  
   197  // Download is called to perform the download.
   198  func (d *mockDownloadDescriptor) Download(ctx context.Context, progressOutput progress.Output) (io.ReadCloser, int64, error) {
   199  	if d.currentDownloads != nil {
   200  		defer atomic.AddInt32(d.currentDownloads, -1)
   201  
   202  		if atomic.AddInt32(d.currentDownloads, 1) > maxDownloadConcurrency {
   203  			return nil, 0, errors.New("concurrency limit exceeded")
   204  		}
   205  	}
   206  
   207  	// Sleep a bit to simulate a time-consuming download.
   208  	for i := int64(0); i <= 10; i++ {
   209  		select {
   210  		case <-ctx.Done():
   211  			return nil, 0, ctx.Err()
   212  		case <-time.After(10 * time.Millisecond):
   213  			progressOutput.WriteProgress(progress.Progress{ID: d.ID(), Action: "Downloading", Current: i, Total: 10})
   214  		}
   215  	}
   216  
   217  	if d.simulateRetries != 0 {
   218  		d.simulateRetries--
   219  		return nil, 0, errors.New("simulating retry")
   220  	}
   221  
   222  	return d.mockTarStream(), 0, nil
   223  }
   224  
   225  func (d *mockDownloadDescriptor) Close() {
   226  }
   227  
   228  func downloadDescriptors(currentDownloads *int32) []DownloadDescriptor {
   229  	return []DownloadDescriptor{
   230  		&mockDownloadDescriptor{
   231  			currentDownloads: currentDownloads,
   232  			id:               "id1",
   233  			expectedDiffID:   layer.DiffID("sha256:68e2c75dc5c78ea9240689c60d7599766c213ae210434c53af18470ae8c53ec1"),
   234  		},
   235  		&mockDownloadDescriptor{
   236  			currentDownloads: currentDownloads,
   237  			id:               "id2",
   238  			expectedDiffID:   layer.DiffID("sha256:64a636223116aa837973a5d9c2bdd17d9b204e4f95ac423e20e65dfbb3655473"),
   239  		},
   240  		&mockDownloadDescriptor{
   241  			currentDownloads: currentDownloads,
   242  			id:               "id3",
   243  			expectedDiffID:   layer.DiffID("sha256:58745a8bbd669c25213e9de578c4da5c8ee1c836b3581432c2b50e38a6753300"),
   244  		},
   245  		&mockDownloadDescriptor{
   246  			currentDownloads: currentDownloads,
   247  			id:               "id2",
   248  			expectedDiffID:   layer.DiffID("sha256:64a636223116aa837973a5d9c2bdd17d9b204e4f95ac423e20e65dfbb3655473"),
   249  		},
   250  		&mockDownloadDescriptor{
   251  			currentDownloads: currentDownloads,
   252  			id:               "id4",
   253  			expectedDiffID:   layer.DiffID("sha256:0dfb5b9577716cc173e95af7c10289322c29a6453a1718addc00c0c5b1330936"),
   254  			simulateRetries:  1,
   255  		},
   256  		&mockDownloadDescriptor{
   257  			currentDownloads: currentDownloads,
   258  			id:               "id5",
   259  			expectedDiffID:   layer.DiffID("sha256:0a5f25fa1acbc647f6112a6276735d0fa01e4ee2aa7ec33015e337350e1ea23d"),
   260  		},
   261  	}
   262  }
   263  
   264  func TestSuccessfulDownload(t *testing.T) {
   265  	// TODO Windows: Fix this unit text
   266  	if runtime.GOOS == "windows" {
   267  		t.Skip("Needs fixing on Windows")
   268  	}
   269  
   270  	layerStore := &mockLayerStore{make(map[layer.ChainID]*mockLayer)}
   271  	lsMap := make(map[string]layer.Store)
   272  	lsMap[runtime.GOOS] = layerStore
   273  	ldm := NewLayerDownloadManager(lsMap, maxDownloadConcurrency, func(m *LayerDownloadManager) { m.waitDuration = time.Millisecond })
   274  
   275  	progressChan := make(chan progress.Progress)
   276  	progressDone := make(chan struct{})
   277  	receivedProgress := make(map[string]progress.Progress)
   278  
   279  	go func() {
   280  		for p := range progressChan {
   281  			receivedProgress[p.ID] = p
   282  		}
   283  		close(progressDone)
   284  	}()
   285  
   286  	var currentDownloads int32
   287  	descriptors := downloadDescriptors(&currentDownloads)
   288  
   289  	firstDescriptor := descriptors[0].(*mockDownloadDescriptor)
   290  
   291  	// Pre-register the first layer to simulate an already-existing layer
   292  	l, err := layerStore.Register(firstDescriptor.mockTarStream(), "")
   293  	if err != nil {
   294  		t.Fatal(err)
   295  	}
   296  	firstDescriptor.diffID = l.DiffID()
   297  
   298  	rootFS, releaseFunc, err := ldm.Download(context.Background(), *image.NewRootFS(), runtime.GOOS, descriptors, progress.ChanOutput(progressChan))
   299  	if err != nil {
   300  		t.Fatalf("download error: %v", err)
   301  	}
   302  
   303  	releaseFunc()
   304  
   305  	close(progressChan)
   306  	<-progressDone
   307  
   308  	if len(rootFS.DiffIDs) != len(descriptors) {
   309  		t.Fatal("got wrong number of diffIDs in rootfs")
   310  	}
   311  
   312  	for i, d := range descriptors {
   313  		descriptor := d.(*mockDownloadDescriptor)
   314  
   315  		if descriptor.diffID != "" {
   316  			if receivedProgress[d.ID()].Action != "Already exists" {
   317  				t.Fatalf("did not get 'Already exists' message for %v", d.ID())
   318  			}
   319  		} else if receivedProgress[d.ID()].Action != "Pull complete" {
   320  			t.Fatalf("did not get 'Pull complete' message for %v", d.ID())
   321  		}
   322  
   323  		if rootFS.DiffIDs[i] != descriptor.expectedDiffID {
   324  			t.Fatalf("rootFS item %d has the wrong diffID (expected: %v got: %v)", i, descriptor.expectedDiffID, rootFS.DiffIDs[i])
   325  		}
   326  
   327  		if descriptor.diffID == "" && descriptor.registeredDiffID != rootFS.DiffIDs[i] {
   328  			t.Fatal("diffID mismatch between rootFS and Registered callback")
   329  		}
   330  	}
   331  }
   332  
   333  func TestCancelledDownload(t *testing.T) {
   334  	layerStore := &mockLayerStore{make(map[layer.ChainID]*mockLayer)}
   335  	lsMap := make(map[string]layer.Store)
   336  	lsMap[runtime.GOOS] = layerStore
   337  	ldm := NewLayerDownloadManager(lsMap, maxDownloadConcurrency, func(m *LayerDownloadManager) { m.waitDuration = time.Millisecond })
   338  	progressChan := make(chan progress.Progress)
   339  	progressDone := make(chan struct{})
   340  
   341  	go func() {
   342  		for range progressChan {
   343  		}
   344  		close(progressDone)
   345  	}()
   346  
   347  	ctx, cancel := context.WithCancel(context.Background())
   348  
   349  	go func() {
   350  		<-time.After(time.Millisecond)
   351  		cancel()
   352  	}()
   353  
   354  	descriptors := downloadDescriptors(nil)
   355  	_, _, err := ldm.Download(ctx, *image.NewRootFS(), runtime.GOOS, descriptors, progress.ChanOutput(progressChan))
   356  	if err != context.Canceled {
   357  		t.Fatal("expected download to be cancelled")
   358  	}
   359  
   360  	close(progressChan)
   361  	<-progressDone
   362  }