Merge pull request #91 from direnv/feature/cache

feature: cache

closes #81

CHANGELOG.md

@@ -5,6 +5,10 @@ All notable changes to the "direnv" extension will be documented in this file.
 Check [Keep a Changelog](http://keepachangelog.com/) for recommendations on how to structure this file.
 
 ## [Unreleased]
+### Added
+- Count added variables separate from changed variables
+- Cache environment variables for faster startup
+- Offer to restart the extension host after receiving updates
 
 ## [0.5.0] - 2021-12-28
 ### Fixed

src/direnv.ts

@@ -39,15 +39,15 @@ function isCommandNotFound(e: unknown, path: string): boolean {
 	return e['path'] === path && e['code'] === 'ENOENT'
 }
 
-const echo: Data = {
+const echo = {
 	['EDITOR']: 'echo',
 }
 
 function cwd() {
 	return vscode.workspace.workspaceFolders?.[0].uri.path ?? process.cwd()
 }
 
-async function direnv(args: string[], env: Data | null = null): Promise<Stdio> {
+async function direnv(args: string[], env?: NodeJS.ProcessEnv): Promise<Stdio> {
 	const options: cp.ExecOptionsWithStringEncoding = {
 		encoding: 'utf8',
 		cwd: cwd(), // same as default cwd for shell tasks
@@ -104,9 +104,7 @@ export async function find(): Promise<string> {
 export async function dump(): Promise<Data> {
 	try {
 		const { stdout } = await direnv(['export', 'json'])
-		if (!stdout) {
-			return {}
-		}
+		if (!stdout) return {}
 		return JSON.parse(stdout) as Data
 	} catch (e) {
 		if (isStdio(e)) {

src/extension.ts

@@ -3,28 +3,42 @@ import * as vscode from 'vscode'
 import * as command from './command'
 import * as config from './config'
 import * as direnv from './direnv'
+import { Data } from './direnv'
 import * as status from './status'
 
+const enum Cached {
+	environment = 'direnv.environment',
+}
+
+const installationUri = vscode.Uri.parse('https://direnv.net/docs/installation.html')
+
 class Direnv implements vscode.Disposable {
 	private backup = new Map<string, string | undefined>()
 	private willLoad = new vscode.EventEmitter<void>()
-	private didLoad = new vscode.EventEmitter<direnv.Data>()
+	private didLoad = new vscode.EventEmitter<Data>()
 	private loaded = new vscode.EventEmitter<void>()
 	private failed = new vscode.EventEmitter<unknown>()
 	private blocked = new vscode.EventEmitter<string>()
 	private viewBlocked = new vscode.EventEmitter<string>()
-	private blockedPath: string | undefined
+	private didUpdate = new vscode.EventEmitter<void>()
+	private blockedPath?: string
 
-	constructor(
-		private environment: vscode.EnvironmentVariableCollection,
-		private status: status.Item,
-	) {
+	constructor(private context: vscode.ExtensionContext, private status: status.Item) {
 		this.willLoad.event(() => this.onWillLoad())
 		this.didLoad.event((e) => this.onDidLoad(e))
 		this.loaded.event(() => this.onLoaded())
 		this.failed.event((e) => this.onFailed(e))
 		this.blocked.event((e) => this.onBlocked(e))
 		this.viewBlocked.event((e) => this.onViewBlocked(e))
+		this.didUpdate.event(() => this.onDidUpdate())
+	}
+
+	private get environment(): vscode.EnvironmentVariableCollection {
+		return this.context.environmentVariableCollection
+	}
+
+	private get cache(): vscode.Memento {
+		return this.context.workspaceState
 	}
 
 	dispose() {
@@ -72,7 +86,7 @@ class Direnv implements vscode.Disposable {
 		await this.open(await direnv.create())
 	}
 
-	async open(path?: string | undefined): Promise<void> {
+	async open(path?: string): Promise<void> {
 		path ??= await direnv.find()
 		const uri = await uriFor(path)
 		const doc = await vscode.workspace.openTextDocument(uri)
@@ -81,13 +95,29 @@ class Direnv implements vscode.Disposable {
 	}
 
 	async reload() {
+		await this.cache.update(Cached.environment, undefined)
 		await this.try(async () => {
 			await direnv.test()
 			this.willLoad.fire()
 		})
 	}
 
-	private updateEnvironment(data: direnv.Data) {
+	restore() {
+		const data = this.cache.get<Data>(Cached.environment)
+		if (data === undefined) return
+		this.updateEnvironment(data)
+	}
+
+	private async updateCache() {
+		await this.cache.update(
+			Cached.environment,
+			Object.fromEntries(
+				[...this.backup.entries()].map(([key]) => [key, process.env[key]]),
+			),
+		)
+	}
+
+	private updateEnvironment(data: Data) {
 		Object.entries(data).forEach(([key, value]) => {
 			if (!this.backup.has(key)) {
 				// keep the oldest value
@@ -99,12 +129,12 @@ class Direnv implements vscode.Disposable {
 				this.environment.replace(key, value)
 			} else {
 				delete process.env[key]
-				this.environment.delete(key)
+				this.environment.delete(key) // can't unset the variable
 			}
 		})
 	}
 
-	private resetEnvironment() {
+	private async resetEnvironment() {
 		this.backup.forEach((value, key) => {
 			if (value === undefined) {
 				delete process.env[key]
@@ -114,6 +144,7 @@ class Direnv implements vscode.Disposable {
 		})
 		this.backup.clear()
 		this.environment.clear()
+		await this.cache.update(Cached.environment, undefined)
 	}
 
 	private async try<T>(callback: () => Promise<T>): Promise<void> {
@@ -137,27 +168,33 @@ class Direnv implements vscode.Disposable {
 		}
 	}
 
-	private onDidLoad(data: direnv.Data) {
+	private async onDidLoad(data: Data) {
 		this.updateEnvironment(data)
+		await this.updateCache()
 		this.loaded.fire()
+		if (Object.keys(data).every(isInternal)) return
+		this.didUpdate.fire()
 	}
 
 	private onLoaded() {
 		let state = status.State.empty
 		if (this.backup.size) {
+			let added = 0
 			let changed = 0
 			let removed = 0
-			this.backup.forEach((_, key) => {
-				if (key in process.env) {
+			this.backup.forEach((value, key) => {
+				if (isInternal(key)) return
+				if (value === undefined) {
+					added += 1
+				} else if (key in process.env) {
 					changed += 1
 				} else {
 					removed += 1
 				}
 			})
-			state = status.State.loaded({ changed, removed })
+			state = status.State.loaded({ added, changed, removed })
 		}
 		this.status.update(state)
-		// TODO: restart extension host here?
 	}
 
 	private async onFailed(err: unknown) {
@@ -169,9 +206,7 @@ class Direnv implements vscode.Disposable {
 				...options,
 			)
 			if (choice === 'Install') {
-				await vscode.env.openExternal(
-					vscode.Uri.parse('https://direnv.net/docs/installation.html'),
-				)
+				await vscode.env.openExternal(installationUri)
 			}
 			if (choice === 'Configure') {
 				await config.path.executable.open()
@@ -186,7 +221,7 @@ class Direnv implements vscode.Disposable {
 
 	private async onBlocked(path: string) {
 		this.blockedPath = path
-		this.resetEnvironment()
+		await this.resetEnvironment()
 		this.status.update(status.State.blocked(path))
 		const options = ['Allow', 'View']
 		const choice = await vscode.window.showWarningMessage(
@@ -210,6 +245,20 @@ class Direnv implements vscode.Disposable {
 			await this.allow(path)
 		}
 	}
+
+	private async onDidUpdate() {
+		const choice = await vscode.window.showWarningMessage(
+			`direnv: Environment updated. Restart extensions?`,
+			'Restart',
+		)
+		if (choice === 'Restart') {
+			await vscode.commands.executeCommand('workbench.action.restartExtensionHost')
+		}
+	}
+}
+
+function isInternal(key: string): boolean {
+	return key.startsWith('DIRENV_')
 }
 
 function message(err: unknown): string | undefined {
@@ -233,9 +282,8 @@ async function uriFor(path: string): Promise<vscode.Uri> {
 }
 
 export async function activate(context: vscode.ExtensionContext) {
-	const environment = context.environmentVariableCollection
 	const statusItem = new status.Item(vscode.window.createStatusBarItem())
-	const instance = new Direnv(environment, statusItem)
+	const instance = new Direnv(context, statusItem)
 	context.subscriptions.push(instance)
 	context.subscriptions.push(
 		vscode.commands.registerCommand(command.Direnv.reload, async () => {
@@ -269,6 +317,7 @@ export async function activate(context: vscode.ExtensionContext) {
 			await instance.configurationChanged(e)
 		}),
 	)
+	instance.restore()
 	await instance.reload()
 }
 

src/status.ts

@@ -3,6 +3,7 @@ import * as command from './command'
 import * as config from './config'
 
 export type Delta = {
+	added: number
 	changed: number
 	removed: number
 }
@@ -11,7 +12,7 @@ export class State {
 	private constructor(
 		readonly text: string,
 		readonly tooltip: string,
-		readonly command: command.Direnv | undefined = undefined,
+		readonly command?: command.Direnv,
 		readonly refresh: () => State = () => this,
 	) {}
 
@@ -20,11 +21,11 @@ export class State {
 	static loaded(delta: Delta): State {
 		let text = '$(folder-active)'
 		if (config.status.showChangesCount.get()) {
-			text += ` +${delta.changed}/-${delta.removed}`
+			text += ` +${delta.added}/~${delta.changed}/-${delta.removed}`
 		}
 		return new State(
 			text,
-			`direnv loaded: ${delta.changed} changed, ${delta.removed} removed\nReload…`,
+			`direnv loaded: ${delta.added} added, ${delta.changed} changed, ${delta.removed} removed\nReload…`,
 			command.Direnv.reload,
 			() => State.loaded(delta),
 		)

src/test/suite/direnv.test.ts

@@ -33,6 +33,7 @@ describe('direnv', () => {
 		})
 
 		it('dumps the allowed .envrc file', async () => {
+			delete process.env['VARIABLE']
 			await direnv.allow(file)
 			const data = await direnv.dump()
 			assert.equal(data['VARIABLE'], 'value')

src/test/suite/extension.test.ts

@@ -54,14 +54,14 @@ describe('the extension', () => {
 
 			describe('with simulated user interaction', () => {
 				it('allows and loads the .envrc on direnv.reload', async () => {
-					sinon.stub(vscode.window, 'showWarningMessage').resolvesArg(1)
+					sinon.stub(vscode.window, 'showWarningMessage').withArgs(sinon.match.any, sinon.match('Allow')).resolvesArg(1)
 					await vscode.commands.executeCommand('direnv.reload')
 					await assertEnvironmentIsLoaded()
 				})
 
 				it('opens and allows and loads the .envrc on direnv.reload', async () => {
-					sinon.stub(vscode.window, 'showWarningMessage').resolvesArg(2)
-					sinon.stub(vscode.window, 'showInformationMessage').resolvesArg(1)
+					sinon.stub(vscode.window, 'showWarningMessage').withArgs(sinon.match.any, sinon.match('Allow')).resolvesArg(2)
+					sinon.stub(vscode.window, 'showInformationMessage').withArgs(sinon.match.any, sinon.match('Allow')).resolvesArg(1)
 					await vscode.commands.executeCommand('direnv.reload')
 					await assertEnvironmentIsLoaded()
 				})