1// Copyright 2017 Google LLC
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
15// Package uid supports generating unique IDs. Its chief purpose is to prevent
16// multiple test executions from interfering with each other, and to facilitate
17// cleanup of old entities that may remain if tests exit early.
18package uid
19
20import (
21	"fmt"
22	"regexp"
23	"strconv"
24	"sync/atomic"
25	"time"
26)
27
28// A Space manages a set of unique IDs distinguished by a prefix.
29type Space struct {
30	Prefix string    // Prefix of UIDs. Read-only.
31	Sep    rune      // Separates UID parts. Read-only.
32	Time   time.Time // Timestamp for UIDs. Read-only.
33	re     *regexp.Regexp
34	count  int32 // atomic
35	short  bool
36}
37
38// Options are optional values for a Space.
39type Options struct {
40	Sep  rune      // Separates parts of the UID. Defaults to '-'.
41	Time time.Time // Timestamp for all UIDs made with this space. Defaults to current time.
42
43	// Short, if true, makes the result of space.New shorter by 6 characters.
44	// This can be useful for character restricted IDs. It will use a shorter
45	// but less readable time representation, and will only use two characters
46	// for the count suffix instead of four.
47	//
48	// e.x. normal: gotest-20181030-59751273685000-0001
49	// e.x. short:  gotest-1540917351273685000-01
50	Short bool
51}
52
53// NewSpace creates a new UID space. A UID Space is used to generate unique IDs.
54func NewSpace(prefix string, opts *Options) *Space {
55	var short bool
56	sep := '-'
57	tm := time.Now().UTC()
58	if opts != nil {
59		short = opts.Short
60		if opts.Sep != 0 {
61			sep = opts.Sep
62		}
63		if !opts.Time.IsZero() {
64			tm = opts.Time
65		}
66	}
67	var re string
68
69	if short {
70		re = fmt.Sprintf(`^%s%[2]c(\d+)%[2]c\d+$`, regexp.QuoteMeta(prefix), sep)
71	} else {
72		re = fmt.Sprintf(`^%s%[2]c(\d{4})(\d{2})(\d{2})%[2]c(\d+)%[2]c\d+$`,
73			regexp.QuoteMeta(prefix), sep)
74	}
75
76	return &Space{
77		Prefix: prefix,
78		Sep:    sep,
79		Time:   tm,
80		re:     regexp.MustCompile(re),
81		short:  short,
82	}
83}
84
85// New generates a new unique ID. The ID consists of the Space's prefix, a
86// timestamp, and a counter value. All unique IDs generated in the same test
87// execution will have the same timestamp.
88//
89// Aside from the characters in the prefix, IDs contain only letters, numbers
90// and sep.
91func (s *Space) New() string {
92	c := atomic.AddInt32(&s.count, 1)
93
94	if s.short && c > 99 {
95		// Short spaces only have space for 99 IDs. (two characters)
96		panic("Short space called New more than 99 times. Ran out of IDs.")
97	} else if c > 9999 {
98		// Spaces only have space for 9999 IDs. (four characters)
99		panic("New called more than 9999 times. Ran out of IDs.")
100	}
101
102	if s.short {
103		return fmt.Sprintf("%s%c%d%c%02d", s.Prefix, s.Sep, s.Time.UnixNano(), s.Sep, c)
104	}
105
106	// Write the time as a date followed by nanoseconds from midnight of that date.
107	// That makes it easier to see the approximate time of the ID when it is displayed.
108	y, m, d := s.Time.Date()
109	ns := s.Time.Sub(time.Date(y, m, d, 0, 0, 0, 0, time.UTC))
110	// Zero-pad the counter for lexical sort order for IDs with the same timestamp.
111	return fmt.Sprintf("%s%c%04d%02d%02d%c%d%c%04d",
112		s.Prefix, s.Sep, y, m, d, s.Sep, ns, s.Sep, c)
113}
114
115// Timestamp extracts the timestamp of uid, which must have been generated by
116// s. The second return value is true on success, false if there was a problem.
117func (s *Space) Timestamp(uid string) (time.Time, bool) {
118	subs := s.re.FindStringSubmatch(uid)
119	if subs == nil {
120		return time.Time{}, false
121	}
122
123	if s.short {
124		ns, err := strconv.ParseInt(subs[1], 10, 64)
125		if err != nil {
126			return time.Time{}, false
127		}
128		return time.Unix(ns/1e9, ns%1e9), true
129	}
130
131	y, err1 := strconv.Atoi(subs[1])
132	m, err2 := strconv.Atoi(subs[2])
133	d, err3 := strconv.Atoi(subs[3])
134	ns, err4 := strconv.Atoi(subs[4])
135	if err1 != nil || err2 != nil || err3 != nil || err4 != nil {
136		return time.Time{}, false
137	}
138	return time.Date(y, time.Month(m), d, 0, 0, 0, ns, time.UTC), true
139}
140
141// Older reports whether uid was created by m and has a timestamp older than
142// the current time by at least d.
143func (s *Space) Older(uid string, d time.Duration) bool {
144	ts, ok := s.Timestamp(uid)
145	if !ok {
146		return false
147	}
148	return time.Since(ts) > d
149}
150