github.com/demonoid81/moby@v0.0.0-20200517203328-62dd8e17c460/distribution/xfer/download_test.go (about)

     1  package xfer // import "github.com/demonoid81/moby/distribution/xfer"
     2  
     3  import (
     4  	"bytes"
     5  	"context"
     6  	"errors"
     7  	"fmt"
     8  	"io"
     9  	"io/ioutil"
    10  	"runtime"
    11  	"sync/atomic"
    12  	"testing"
    13  	"time"
    14  
    15  	"github.com/docker/distribution"
    16  	"github.com/demonoid81/moby/image"
    17  	"github.com/demonoid81/moby/layer"
    18  	"github.com/demonoid81/moby/pkg/progress"
    19  	digest "github.com/opencontainers/go-digest"
    20  	"gotest.tools/v3/assert"
    21  )
    22  
    23  const maxDownloadConcurrency = 3
    24  
    25  type mockLayer struct {
    26  	layerData bytes.Buffer
    27  	diffID    layer.DiffID
    28  	chainID   layer.ChainID
    29  	parent    layer.Layer
    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  	retries          int
   165  }
   166  
   167  // Key returns the key used to deduplicate downloads.
   168  func (d *mockDownloadDescriptor) Key() string {
   169  	return d.id
   170  }
   171  
   172  // ID returns the ID for display purposes.
   173  func (d *mockDownloadDescriptor) ID() string {
   174  	return d.id
   175  }
   176  
   177  // DiffID should return the DiffID for this layer, or an error
   178  // if it is unknown (for example, if it has not been downloaded
   179  // before).
   180  func (d *mockDownloadDescriptor) DiffID() (layer.DiffID, error) {
   181  	if d.diffID != "" {
   182  		return d.diffID, nil
   183  	}
   184  	return "", errors.New("no diffID available")
   185  }
   186  
   187  func (d *mockDownloadDescriptor) Registered(diffID layer.DiffID) {
   188  	d.registeredDiffID = diffID
   189  }
   190  
   191  func (d *mockDownloadDescriptor) mockTarStream() io.ReadCloser {
   192  	// The mock implementation returns the ID repeated 5 times as a tar
   193  	// stream instead of actual tar data. The data is ignored except for
   194  	// computing IDs.
   195  	return ioutil.NopCloser(bytes.NewBuffer([]byte(d.id + d.id + d.id + d.id + d.id)))
   196  }
   197  
   198  // Download is called to perform the download.
   199  func (d *mockDownloadDescriptor) Download(ctx context.Context, progressOutput progress.Output) (io.ReadCloser, int64, error) {
   200  	if d.currentDownloads != nil {
   201  		defer atomic.AddInt32(d.currentDownloads, -1)
   202  
   203  		if atomic.AddInt32(d.currentDownloads, 1) > maxDownloadConcurrency {
   204  			return nil, 0, errors.New("concurrency limit exceeded")
   205  		}
   206  	}
   207  
   208  	// Sleep a bit to simulate a time-consuming download.
   209  	for i := int64(0); i <= 10; i++ {
   210  		select {
   211  		case <-ctx.Done():
   212  			return nil, 0, ctx.Err()
   213  		case <-time.After(10 * time.Millisecond):
   214  			progressOutput.WriteProgress(progress.Progress{ID: d.ID(), Action: "Downloading", Current: i, Total: 10})
   215  		}
   216  	}
   217  
   218  	if d.retries < d.simulateRetries {
   219  		d.retries++
   220  		return nil, 0, fmt.Errorf("simulating download attempt %d/%d", d.retries, d.simulateRetries)
   221  	}
   222  
   223  	return d.mockTarStream(), 0, nil
   224  }
   225  
   226  func (d *mockDownloadDescriptor) Close() {
   227  }
   228  
   229  func downloadDescriptors(currentDownloads *int32) []DownloadDescriptor {
   230  	return []DownloadDescriptor{
   231  		&mockDownloadDescriptor{
   232  			currentDownloads: currentDownloads,
   233  			id:               "id1",
   234  			expectedDiffID:   layer.DiffID("sha256:68e2c75dc5c78ea9240689c60d7599766c213ae210434c53af18470ae8c53ec1"),
   235  		},
   236  		&mockDownloadDescriptor{
   237  			currentDownloads: currentDownloads,
   238  			id:               "id2",
   239  			expectedDiffID:   layer.DiffID("sha256:64a636223116aa837973a5d9c2bdd17d9b204e4f95ac423e20e65dfbb3655473"),
   240  		},
   241  		&mockDownloadDescriptor{
   242  			currentDownloads: currentDownloads,
   243  			id:               "id3",
   244  			expectedDiffID:   layer.DiffID("sha256:58745a8bbd669c25213e9de578c4da5c8ee1c836b3581432c2b50e38a6753300"),
   245  		},
   246  		&mockDownloadDescriptor{
   247  			currentDownloads: currentDownloads,
   248  			id:               "id2",
   249  			expectedDiffID:   layer.DiffID("sha256:64a636223116aa837973a5d9c2bdd17d9b204e4f95ac423e20e65dfbb3655473"),
   250  		},
   251  		&mockDownloadDescriptor{
   252  			currentDownloads: currentDownloads,
   253  			id:               "id4",
   254  			expectedDiffID:   layer.DiffID("sha256:0dfb5b9577716cc173e95af7c10289322c29a6453a1718addc00c0c5b1330936"),
   255  			simulateRetries:  1,
   256  		},
   257  		&mockDownloadDescriptor{
   258  			currentDownloads: currentDownloads,
   259  			id:               "id5",
   260  			expectedDiffID:   layer.DiffID("sha256:0a5f25fa1acbc647f6112a6276735d0fa01e4ee2aa7ec33015e337350e1ea23d"),
   261  		},
   262  	}
   263  }
   264  
   265  func TestSuccessfulDownload(t *testing.T) {
   266  	// TODO Windows: Fix this unit text
   267  	if runtime.GOOS == "windows" {
   268  		t.Skip("Needs fixing on Windows")
   269  	}
   270  
   271  	layerStore := &mockLayerStore{make(map[layer.ChainID]*mockLayer)}
   272  	lsMap := make(map[string]layer.Store)
   273  	lsMap[runtime.GOOS] = layerStore
   274  	ldm := NewLayerDownloadManager(lsMap, maxDownloadConcurrency, func(m *LayerDownloadManager) { m.waitDuration = time.Millisecond })
   275  
   276  	progressChan := make(chan progress.Progress)
   277  	progressDone := make(chan struct{})
   278  	receivedProgress := make(map[string]progress.Progress)
   279  
   280  	go func() {
   281  		for p := range progressChan {
   282  			receivedProgress[p.ID] = p
   283  		}
   284  		close(progressDone)
   285  	}()
   286  
   287  	var currentDownloads int32
   288  	descriptors := downloadDescriptors(&currentDownloads)
   289  
   290  	firstDescriptor := descriptors[0].(*mockDownloadDescriptor)
   291  
   292  	// Pre-register the first layer to simulate an already-existing layer
   293  	l, err := layerStore.Register(firstDescriptor.mockTarStream(), "")
   294  	if err != nil {
   295  		t.Fatal(err)
   296  	}
   297  	firstDescriptor.diffID = l.DiffID()
   298  
   299  	rootFS, releaseFunc, err := ldm.Download(context.Background(), *image.NewRootFS(), runtime.GOOS, descriptors, progress.ChanOutput(progressChan))
   300  	if err != nil {
   301  		t.Fatalf("download error: %v", err)
   302  	}
   303  
   304  	releaseFunc()
   305  
   306  	close(progressChan)
   307  	<-progressDone
   308  
   309  	if len(rootFS.DiffIDs) != len(descriptors) {
   310  		t.Fatal("got wrong number of diffIDs in rootfs")
   311  	}
   312  
   313  	for i, d := range descriptors {
   314  		descriptor := d.(*mockDownloadDescriptor)
   315  
   316  		if descriptor.diffID != "" {
   317  			if receivedProgress[d.ID()].Action != "Already exists" {
   318  				t.Fatalf("did not get 'Already exists' message for %v", d.ID())
   319  			}
   320  		} else if receivedProgress[d.ID()].Action != "Pull complete" {
   321  			t.Fatalf("did not get 'Pull complete' message for %v", d.ID())
   322  		}
   323  
   324  		if rootFS.DiffIDs[i] != descriptor.expectedDiffID {
   325  			t.Fatalf("rootFS item %d has the wrong diffID (expected: %v got: %v)", i, descriptor.expectedDiffID, rootFS.DiffIDs[i])
   326  		}
   327  
   328  		if descriptor.diffID == "" && descriptor.registeredDiffID != rootFS.DiffIDs[i] {
   329  			t.Fatal("diffID mismatch between rootFS and Registered callback")
   330  		}
   331  	}
   332  }
   333  
   334  func TestCancelledDownload(t *testing.T) {
   335  	layerStore := &mockLayerStore{make(map[layer.ChainID]*mockLayer)}
   336  	lsMap := make(map[string]layer.Store)
   337  	lsMap[runtime.GOOS] = layerStore
   338  	ldm := NewLayerDownloadManager(lsMap, maxDownloadConcurrency, func(m *LayerDownloadManager) { m.waitDuration = time.Millisecond })
   339  	progressChan := make(chan progress.Progress)
   340  	progressDone := make(chan struct{})
   341  
   342  	go func() {
   343  		for range progressChan {
   344  		}
   345  		close(progressDone)
   346  	}()
   347  
   348  	ctx, cancel := context.WithCancel(context.Background())
   349  
   350  	go func() {
   351  		<-time.After(time.Millisecond)
   352  		cancel()
   353  	}()
   354  
   355  	descriptors := downloadDescriptors(nil)
   356  	_, _, err := ldm.Download(ctx, *image.NewRootFS(), runtime.GOOS, descriptors, progress.ChanOutput(progressChan))
   357  	if err != context.Canceled {
   358  		t.Fatal("expected download to be cancelled")
   359  	}
   360  
   361  	close(progressChan)
   362  	<-progressDone
   363  }
   364  
   365  func TestMaxDownloadAttempts(t *testing.T) {
   366  	tests := []struct {
   367  		name                string
   368  		simulateRetries     int
   369  		maxDownloadAttempts int
   370  		expectedErr         string
   371  	}{
   372  		{
   373  			name:                "max-attempts=5, succeed at 2nd attempt",
   374  			simulateRetries:     2,
   375  			maxDownloadAttempts: 5,
   376  		},
   377  		{
   378  			name:                "max-attempts=5, succeed at 5th attempt",
   379  			simulateRetries:     5,
   380  			maxDownloadAttempts: 5,
   381  		},
   382  		{
   383  			name:                "max-attempts=5, fail at 6th attempt",
   384  			simulateRetries:     6,
   385  			maxDownloadAttempts: 5,
   386  			expectedErr:         "simulating download attempt 5/6",
   387  		},
   388  		{
   389  			name:                "max-attempts=0, fail after 1 attempt",
   390  			simulateRetries:     1,
   391  			maxDownloadAttempts: 0,
   392  			expectedErr:         "simulating download attempt 1/1",
   393  		},
   394  	}
   395  	for _, tc := range tests {
   396  		t.Run(tc.name, func(t *testing.T) {
   397  			t.Parallel()
   398  			layerStore := &mockLayerStore{make(map[layer.ChainID]*mockLayer)}
   399  			lsMap := make(map[string]layer.Store)
   400  			lsMap[runtime.GOOS] = layerStore
   401  			ldm := NewLayerDownloadManager(
   402  				lsMap,
   403  				maxDownloadConcurrency,
   404  				func(m *LayerDownloadManager) {
   405  					m.waitDuration = time.Millisecond
   406  					m.maxDownloadAttempts = tc.maxDownloadAttempts
   407  				})
   408  
   409  			progressChan := make(chan progress.Progress)
   410  			progressDone := make(chan struct{})
   411  
   412  			go func() {
   413  				for range progressChan {
   414  				}
   415  				close(progressDone)
   416  			}()
   417  
   418  			var currentDownloads int32
   419  			descriptors := downloadDescriptors(&currentDownloads)
   420  			descriptors[4].(*mockDownloadDescriptor).simulateRetries = tc.simulateRetries
   421  
   422  			_, _, err := ldm.Download(context.Background(), *image.NewRootFS(), runtime.GOOS, descriptors, progress.ChanOutput(progressChan))
   423  			if tc.expectedErr == "" {
   424  				assert.NilError(t, err)
   425  			} else {
   426  				assert.Error(t, err, tc.expectedErr)
   427  			}
   428  		})
   429  	}
   430  }