1 files changed, 61 insertions, 10 deletions

diff --git a/services/repository/archiver/archiver.go b/services/repository/archiver/archiver.go
index ebd3eaf236..ae43503bae 100644
--- a/services/repository/archiver/archiver.go
+++ b/services/repository/archiver/archiver.go
@@ -57,6 +57,21 @@ func (ErrUnknownArchiveFormat) Is(err error) bool {
 	return ok
 }
 
+// RepoRefNotFoundError is returned when a requested reference (commit, tag) was not found.
+type RepoRefNotFoundError struct {
+	RefName string
+}
+
+// Error implements error.
+func (e RepoRefNotFoundError) Error() string {
+	return fmt.Sprintf("unrecognized repository reference: %s", e.RefName)
+}
+
+func (e RepoRefNotFoundError) Is(err error) bool {
+	_, ok := err.(RepoRefNotFoundError)
+	return ok
+}
+
 // NewRequest creates an archival request, based on the URI.  The
 // resulting ArchiveRequest is suitable for being passed to ArchiveRepository()
 // if it's determined that the request still needs to be satisfied.
@@ -103,7 +118,7 @@ func NewRequest(repoID int64, repo *git.Repository, uri string) (*ArchiveRequest
 			}
 		}
 	} else {
-		return nil, fmt.Errorf("Unknow ref %s type", r.refName)
+		return nil, RepoRefNotFoundError{RefName: r.refName}
 	}
 
 	return r, nil
@@ -115,6 +130,49 @@ func (aReq *ArchiveRequest) GetArchiveName() string {
 	return strings.ReplaceAll(aReq.refName, "/", "-") + "." + aReq.Type.String()
 }
 
+// Await awaits the completion of an ArchiveRequest. If the archive has
+// already been prepared the method returns immediately. Otherwise an archiver
+// process will be started and its completion awaited. On success the returned
+// RepoArchiver may be used to download the archive. Note that even if the
+// context is cancelled/times out a started archiver will still continue to run
+// in the background.
+func (aReq *ArchiveRequest) Await(ctx context.Context) (*repo_model.RepoArchiver, error) {
+	archiver, err := repo_model.GetRepoArchiver(ctx, aReq.RepoID, aReq.Type, aReq.CommitID)
+	if err != nil {
+		return nil, fmt.Errorf("models.GetRepoArchiver: %v", err)
+	}
+
+	if archiver != nil && archiver.Status == repo_model.ArchiverReady {
+		// Archive already generated, we're done.
+		return archiver, nil
+	}
+
+	if err := StartArchive(aReq); err != nil {
+		return nil, fmt.Errorf("archiver.StartArchive: %v", err)
+	}
+
+	poll := time.NewTicker(time.Second * 1)
+	defer poll.Stop()
+
+	for {
+		select {
+		case <-graceful.GetManager().HammerContext().Done():
+			// System stopped.
+			return nil, graceful.GetManager().HammerContext().Err()
+		case <-ctx.Done():
+			return nil, ctx.Err()
+		case <-poll.C:
+			archiver, err = repo_model.GetRepoArchiver(ctx, aReq.RepoID, aReq.Type, aReq.CommitID)
+			if err != nil {
+				return nil, fmt.Errorf("repo_model.GetRepoArchiver: %v", err)
+			}
+			if archiver != nil && archiver.Status == repo_model.ArchiverReady {
+				return archiver, nil
+			}
+		}
+	}
+}
+
 func doArchive(r *ArchiveRequest) (*repo_model.RepoArchiver, error) {
 	txCtx, committer, err := db.TxContext()
 	if err != nil {
@@ -147,11 +205,7 @@ func doArchive(r *ArchiveRequest) (*repo_model.RepoArchiver, error) {
 		}
 	}
 
-	rPath, err := archiver.RelativePath()
-	if err != nil {
-		return nil, err
-	}
-
+	rPath := archiver.RelativePath()
 	_, err = storage.RepoArchives.Stat(rPath)
 	if err == nil {
 		if archiver.Status == repo_model.ArchiverGenerating {
@@ -284,13 +338,10 @@ func StartArchive(request *ArchiveRequest) error {
 }
 
 func deleteOldRepoArchiver(ctx context.Context, archiver *repo_model.RepoArchiver) error {
-	p, err := archiver.RelativePath()
-	if err != nil {
-		return err
-	}
 	if err := repo_model.DeleteRepoArchiver(ctx, archiver); err != nil {
 		return err
 	}
+	p := archiver.RelativePath()
 	if err := storage.RepoArchives.Delete(p); err != nil {
 		log.Error("delete repo archive file failed: %v", err)
 	}